Use ChunkStoreHandle in the OSS server instead of EntityDb (#11422)

Co-authored-by: Clement Rey <[email protected]>

Author: abey79

Cargo.lock

@@ -9384,13 +9384,16 @@ dependencies = [
  "nohash-hasher",
  "re_build_info",
  "re_build_tools",
+ "re_byte_size",
  "re_chunk_store",
  "re_entity_db",
  "re_grpc_server",
  "re_log",
  "re_log_encoding",
  "re_log_types",
  "re_protos",
+ "re_tuid",
+ "re_types_core",
  "thiserror 1.0.69",
  "tokio",
  "tokio-stream",

crates/store/re_server/Cargo.toml

@@ -29,12 +29,15 @@ default = []
 # Rerun
 re_chunk_store.workspace = true
 re_build_info.workspace = true
+re_byte_size.workspace = true
 re_entity_db.workspace = true
 re_grpc_server.workspace = true
 re_log = { workspace = true, features = ["setup"] }
 re_log_encoding.workspace = true
 re_log_types.workspace = true
 re_protos.workspace = true
+re_tuid.workspace = true
+re_types_core.workspace = true
 
 # External
 anyhow.workspace = true

crates/store/re_server/src/rerun_cloud.rs

@@ -9,41 +9,30 @@ use arrow::array::{
 use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
 use datafusion::prelude::SessionContext;
 use nohash_hasher::IntSet;
-use re_chunk_store::Chunk;
-use re_chunk_store::external::re_chunk::external::re_byte_size::SizeBytes as _;
-use re_entity_db::EntityDb;
-use re_entity_db::external::re_query::StorageEngine;
+use tokio_stream::StreamExt as _;
+use tonic::{Code, Status};
+
+use re_byte_size::SizeBytes as _;
+use re_chunk_store::{Chunk, ChunkStore, ChunkStoreConfig, ChunkStoreHandle};
 use re_log_encoding::codec::wire::{decoder::Decode as _, encoder::Encode as _};
-use re_log_types::external::re_types_core::{ChunkId, Loggable as _};
 use re_log_types::{EntityPath, EntryId, StoreId, StoreKind};
-use re_protos::cloud::v1alpha1::ext::GetChunksRequest;
-use re_protos::cloud::v1alpha1::{
-    EntryDetails, GetChunksResponse, GetDatasetSchemaResponse, GetPartitionTableSchemaResponse,
-    QueryDatasetResponse, ScanPartitionTableResponse, ScanTableResponse,
-};
-use re_protos::headers::RerunHeadersExtractorExt as _;
-use re_protos::{cloud::v1alpha1::RegisterWithDatasetResponse, common::v1alpha1::ext::PartitionId};
 use re_protos::{
-    cloud::v1alpha1::ext,
-    cloud::v1alpha1::ext::{
-        CreateDatasetEntryResponse, ReadDatasetEntryResponse, ReadTableEntryResponse,
-    },
-};
-use re_protos::{
-    cloud::v1alpha1::rerun_cloud_service_server::RerunCloudService,
     cloud::v1alpha1::{
-        FetchTaskOutputRequest, FetchTaskOutputResponse, QueryTasksOnCompletionRequest,
-        QueryTasksRequest, QueryTasksResponse,
+        DeleteEntryResponse, EntryDetails, EntryKind, FetchTaskOutputRequest,
+        FetchTaskOutputResponse, GetChunksResponse, GetDatasetSchemaResponse,
+        GetPartitionTableSchemaResponse, QueryDatasetResponse, QueryTasksOnCompletionRequest,
+        QueryTasksRequest, QueryTasksResponse, RegisterTableRequest, RegisterTableResponse,
+        RegisterWithDatasetResponse, ScanPartitionTableResponse, ScanTableResponse,
+        ext::{
+            self, CreateDatasetEntryResponse, GetChunksRequest, ReadDatasetEntryResponse,
+            ReadTableEntryResponse,
+        },
+        rerun_cloud_service_server::RerunCloudService,
     },
+    common::v1alpha1::ext::{IfDuplicateBehavior, PartitionId},
+    headers::RerunHeadersExtractorExt as _,
 };
-use re_protos::{
-    cloud::v1alpha1::{
-        DeleteEntryResponse, EntryKind, RegisterTableRequest, RegisterTableResponse,
-    },
-    common::v1alpha1::ext::IfDuplicateBehavior,
-};
-use tokio_stream::StreamExt as _;
-use tonic::{Code, Status};
+use re_types_core::{ChunkId, Loggable as _};
 
 use crate::store::{Dataset, InMemoryStore, Table};
 
@@ -111,7 +100,7 @@ impl RerunCloudHandler {
         &self,
         dataset_id: EntryId,
         mut partition_ids: Vec<PartitionId>,
-    ) -> Result<Vec<(PartitionId, StorageEngine)>, tonic::Status> {
+    ) -> Result<Vec<(PartitionId, ChunkStoreHandle)>, tonic::Status> {
         let store = self.store.read().await;
         let dataset = store.dataset(dataset_id).ok_or_else(|| {
             tonic::Status::not_found(format!("Entry with ID {dataset_id} not found"))
@@ -125,19 +114,13 @@ impl RerunCloudHandler {
             .into_iter()
             .map(|partition_id| {
                 dataset
-                    .partition(&partition_id)
+                    .partition_store_handle(&partition_id)
                     .ok_or_else(|| {
                         tonic::Status::not_found(format!(
                             "Partition with ID {partition_id} not found"
                         ))
                     })
-                    .map(|partition| {
-                        #[expect(unsafe_code)]
-                        // Safety: no viewer is running, and we've locked the store for the duration
-                        // of the handler already.
-                        unsafe { partition.storage_engine_raw() }.clone()
-                    })
-                    .map(|storage_engine| (partition_id, storage_engine))
+                    .map(|store_handle| (partition_id, store_handle.clone()))
             })
             .collect::<Result<Vec<_>, _>>()
     }
@@ -560,7 +543,7 @@ impl RerunCloudService for RerunCloudHandler {
 
         let mut request = request.into_inner();
 
-        let mut entity_dbs = HashMap::new();
+        let mut chunk_stores = HashMap::new();
 
         while let Some(chunk_msg) = request.next().await {
             let chunk_msg = chunk_msg?;
@@ -589,16 +572,15 @@ impl RerunCloudService for RerunCloudHandler {
                 tonic::Status::internal(format!("error decoding chunk from record batch: {err:#}"))
             })?);
 
-            entity_dbs
+            chunk_stores
                 .entry(partition_id.clone())
                 .or_insert_with(|| {
-                    EntityDb::new(StoreId::new(
-                        StoreKind::Recording,
-                        entry_id.to_string(),
-                        partition_id.id,
-                    ))
+                    ChunkStore::new(
+                        StoreId::new(StoreKind::Recording, entry_id.to_string(), partition_id.id),
+                        ChunkStoreConfig::CHANGELOG_DISABLED,
+                    )
                 })
-                .add_chunk(&chunk)
+                .insert_chunk(&chunk)
                 .map_err(|err| {
                     tonic::Status::internal(format!("error adding chunk to store: {err:#}"))
                 })?;
@@ -610,8 +592,8 @@ impl RerunCloudService for RerunCloudHandler {
         };
 
         #[expect(clippy::iter_over_hash_type)]
-        for (entity_path, entity_db) in entity_dbs {
-            dataset.add_partition(entity_path, entity_db);
+        for (entity_path, chunk_store) in chunk_stores {
+            dataset.add_partition(entity_path, ChunkStoreHandle::new(chunk_store));
         }
 
         Ok(tonic::Response::new(
@@ -767,9 +749,8 @@ impl RerunCloudService for RerunCloudHandler {
         let storage_engines = self.get_storage_engines(entry_id, partition_ids).await?;
 
         let stream = futures::stream::iter(storage_engines.into_iter().map(
-            move |(partition_id, storage_engine)| {
-                let storage_read = storage_engine.read();
-                let chunk_store = storage_read.store();
+            move |(partition_id, store_handle)| {
+                let chunk_store = store_handle.read();
                 let num_rows = chunk_store.num_chunks();
 
                 let mut chunk_partition_id = Vec::with_capacity(num_rows);
@@ -940,21 +921,20 @@ impl RerunCloudService for RerunCloudHandler {
         let storage_engines = self.get_storage_engines(entry_id, partition_ids).await?;
 
         let stream = futures::stream::iter(storage_engines.into_iter().map(
-            move |(partition_id, storage_engine)| {
+            move |(partition_id, store_handle)| {
                 let compression = re_log_encoding::Compression::Off;
                 let store_id = StoreId::new(
                     StoreKind::Recording,
                     entry_id.to_string(),
                     partition_id.id.as_str(),
                 );
 
-                let arrow_msgs: Result<Vec<_>, _> = storage_engine
+                let arrow_msgs: Result<Vec<_>, _> = store_handle
                     // NOTE: ⚠️This is super cursed ⚠️The underlying lock is synchronous: the only
                     // reason this doesn't deadlock is because we collect() at the end of this mapping,
                     // before the overarching stream ever gets a chance to yield.
                     // Make sure it stays that way.
                     .read()
-                    .store()
                     .iter_chunks()
                     .filter(|chunk| {
                         entity_paths.is_empty() || entity_paths.contains(chunk.entity_path())
@@ -1036,7 +1016,7 @@ impl RerunCloudService for RerunCloudHandler {
 
         // get storage engines only for the partitions we actually need
         let store = self.store.read().await;
-        let storage_engines: std::collections::HashMap<_, _> = store
+        let store_handles: std::collections::HashMap<_, _> = store
             .iter_datasets()
             .flat_map(|dataset| {
                 let dataset_id = dataset.id();
@@ -1046,13 +1026,9 @@ impl RerunCloudService for RerunCloudHandler {
                         .iter()
                         .any(|(_, pid)| pid == &partition_id)
                     {
-                        dataset.partition(&partition_id).map(|partition| {
-                            #[expect(unsafe_code)]
-                            // Safety: no viewer is running, and we've locked the store for the duration
-                            // of the handler already.
-                            let storage_engine = unsafe { partition.storage_engine_raw() }.clone();
-                            (partition_id, (dataset_id, storage_engine))
-                        })
+                        dataset
+                            .partition_store_handle(&partition_id)
+                            .map(|store_handle| (partition_id, (dataset_id, store_handle.clone())))
                     } else {
                         None
                     }
@@ -1065,15 +1041,13 @@ impl RerunCloudService for RerunCloudHandler {
         let compression = re_log_encoding::Compression::Off;
 
         for (chunk_id, partition_id) in chunk_partition_pairs {
-            let (dataset_id, storage_engine) =
-                storage_engines.get(&partition_id).ok_or_else(|| {
-                    tonic::Status::internal(format!(
-                        "Storage engine not found for partition {partition_id}"
-                    ))
-                })?;
+            let (dataset_id, store_handle) = store_handles.get(&partition_id).ok_or_else(|| {
+                tonic::Status::internal(format!(
+                    "Storage engine not found for partition {partition_id}"
+                ))
+            })?;
 
-            let storage_read = storage_engine.read();
-            let chunk_store = storage_read.store();
+            let chunk_store = store_handle.read();
 
             if let Some(chunk) = chunk_store.chunk(&chunk_id) {
                 let store_id = StoreId::new(

crates/store/re_server/src/store.rs

@@ -1,3 +1,9 @@
+use std::{
+    collections::{BTreeSet, HashMap, hash_map::Entry},
+    path::Path,
+    sync::Arc,
+};
+
 use arrow::array::{
     ArrayRef, Int32Array, RecordBatchOptions, StringArray, TimestampNanosecondArray,
 };
@@ -9,26 +15,18 @@ use datafusion::error::DataFusionError;
 use itertools::Itertools as _;
 use jiff::Timestamp;
 use lance::datafusion::LanceTableProvider;
-use re_entity_db::{EntityDb, StoreBundle};
-use re_log_types::external::re_tuid::Tuid;
-use re_log_types::external::re_types_core::{ComponentBatch as _, Loggable as _};
+
+use re_chunk_store::{ChunkStore, ChunkStoreConfig, ChunkStoreHandle};
 use re_log_types::{EntryId, StoreKind};
-use re_protos::cloud::v1alpha1::SystemTableKind;
-use re_protos::cloud::v1alpha1::ext::{ProviderDetails as _, SystemTable};
 use re_protos::{
-    cloud::v1alpha1::ScanPartitionTableResponse,
     cloud::v1alpha1::{
-        EntryKind,
-        ext::{DatasetEntry, EntryDetails, TableEntry},
+        EntryKind, ScanPartitionTableResponse, SystemTableKind,
+        ext::{DatasetEntry, EntryDetails, ProviderDetails as _, SystemTable, TableEntry},
     },
     common::v1alpha1::ext::{DatasetHandle, IfDuplicateBehavior, PartitionId},
 };
-use std::sync::Arc;
-use std::{
-    collections::{BTreeSet, HashMap, hash_map::Entry},
-    fs::File,
-    path::Path,
-};
+use re_tuid::Tuid;
+use re_types_core::{ComponentBatch as _, Loggable as _};
 
 const ENTRIES_TABLE_NAME: &str = "__entries";
 
@@ -49,6 +47,9 @@ pub enum Error {
 
     #[error(transparent)]
     DataFusionError(#[from] datafusion::error::DataFusionError),
+
+    #[error("Error loading RRD: {0}")]
+    RrdLoadingError(anyhow::Error),
 }
 
 impl From<Error> for tonic::Status {
@@ -61,12 +62,13 @@ impl From<Error> for tonic::Status {
             }
             Error::EntryIdNotFound(id) => Self::not_found(format!("Entry ID not found: {id}")),
             Error::DataFusionError(err) => Self::internal(format!("DataFusion error: {err:#}")),
+            Error::RrdLoadingError(err) => Self::internal(format!("{err:#}")),
         }
     }
 }
 
 pub struct Partition {
-    entity_db: EntityDb,
+    store_handle: ChunkStoreHandle,
     registration_time: jiff::Timestamp,
 }
 
@@ -116,7 +118,7 @@ impl Dataset {
 
     pub fn schema(&self) -> arrow::error::Result<Schema> {
         let schemas = self.partitions.values().map(|partition| {
-            let columns = partition.entity_db.storage_engine().store().schema();
+            let columns = partition.store_handle.read().schema();
             let fields = columns.arrow_fields();
             Schema::new_with_metadata(fields, HashMap::default())
         });
@@ -160,16 +162,16 @@ impl Dataset {
         )
     }
 
-    pub fn partition(&self, partition_id: &PartitionId) -> Option<&EntityDb> {
-        self.partitions.get(partition_id).map(|p| &p.entity_db)
+    pub fn partition_store_handle(&self, partition_id: &PartitionId) -> Option<&ChunkStoreHandle> {
+        self.partitions.get(partition_id).map(|p| &p.store_handle)
     }
 
-    pub fn add_partition(&mut self, partition_id: PartitionId, entity_db: EntityDb) {
+    pub fn add_partition(&mut self, partition_id: PartitionId, store_handle: ChunkStoreHandle) {
         re_log::debug!(?partition_id, "add_partition");
         self.partitions.insert(
             partition_id,
             Partition {
-                entity_db,
+                store_handle,
                 registration_time: jiff::Timestamp::now(),
             },
         );
@@ -182,12 +184,13 @@ impl Dataset {
         on_duplicate: IfDuplicateBehavior,
     ) -> Result<BTreeSet<PartitionId>, Error> {
         re_log::info!("Loading RRD: {}", path.display());
-        let mut contents = StoreBundle::from_rrd(File::open(path)?)?;
+        let contents =
+            ChunkStore::handle_from_rrd_filepath(&ChunkStoreConfig::CHANGELOG_DISABLED, path)
+                .map_err(Error::RrdLoadingError)?;
 
         let mut new_partition_ids = BTreeSet::default();
 
-        for entity_db in contents.drain_entity_dbs() {
-            let store_id = entity_db.store_id();
+        for (store_id, chunk_store) in contents {
             if !store_id.is_recording() {
                 continue;
             }
@@ -198,15 +201,15 @@ impl Dataset {
                 Entry::Vacant(entry) => {
                     new_partition_ids.insert(partition_id);
                     entry.insert(Partition {
-                        entity_db,
+                        store_handle: chunk_store,
                         registration_time: jiff::Timestamp::now(),
                     });
                 }
                 Entry::Occupied(mut entry) => match on_duplicate {
                     IfDuplicateBehavior::Overwrite => {
                         re_log::info!("Overwriting {partition_id}");
                         entry.insert(Partition {
-                            entity_db,
+                            store_handle: chunk_store,
                             registration_time: jiff::Timestamp::now(),
                         });
                     }