datastore: incremental metadata registry stats

crates/re_arrow_store/src/store.rs

@@ -97,22 +97,38 @@ impl std::ops::DerefMut for DataTypeRegistry {
 }
 
 /// Keeps track of arbitrary per-row metadata.
-#[derive(Debug, Default, Clone)]
-pub struct MetadataRegistry<T: Clone>(pub BTreeMap<RowId, T>);
+#[derive(Debug, Clone)]
+pub struct MetadataRegistry<T: Clone> {
+    pub registry: BTreeMap<RowId, T>,
+
+    /// Cached heap size, because the registry gets very, very large.
+    pub heap_size_bytes: u64,
+}
+
+impl Default for MetadataRegistry<TimePoint> {
+    fn default() -> Self {
+        let mut this = Self {
+            registry: Default::default(),
+            heap_size_bytes: 0,
+        };
+        this.heap_size_bytes = this.heap_size_bytes(); // likely zero, just future proofing
+        this
+    }
+}
 
 impl<T: Clone> std::ops::Deref for MetadataRegistry<T> {
     type Target = BTreeMap<RowId, T>;
 
     #[inline]
     fn deref(&self) -> &Self::Target {
-        &self.0
+        &self.registry
     }
 }
 
 impl<T: Clone> std::ops::DerefMut for MetadataRegistry<T> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
+        &mut self.registry
     }
 }
 

crates/re_arrow_store/src/store_gc.rs

@@ -139,8 +139,10 @@ impl DataStore {
             let Some((row_id, timepoint)) = self.metadata_registry.pop_first() else {
                 break;
             };
-            num_bytes_to_drop -= row_id.total_size_bytes() as f64;
-            num_bytes_to_drop -= timepoint.total_size_bytes() as f64;
+            let metadata_dropped_size_bytes =
+                row_id.total_size_bytes() + timepoint.total_size_bytes();
+            self.metadata_registry.heap_size_bytes -= metadata_dropped_size_bytes;
+            num_bytes_to_drop -= metadata_dropped_size_bytes as f64;
             row_ids.push(row_id);
 
             // find all tables that could possibly contain this `RowId`

crates/re_arrow_store/src/store_stats.rs

@@ -83,31 +83,48 @@ impl DataStoreStats {
     pub fn from_store(store: &DataStore) -> Self {
         crate::profile_function!();
 
-        let type_registry = DataStoreRowStats {
-            num_rows: store.type_registry.len() as _,
-            num_bytes: store.type_registry.total_size_bytes(),
+        let type_registry = {
+            crate::profile_scope!("type_registry");
+            DataStoreRowStats {
+                num_rows: store.type_registry.len() as _,
+                num_bytes: store.type_registry.total_size_bytes(),
+            }
         };
 
-        let metadata_registry = DataStoreRowStats {
-            num_rows: store.metadata_registry.len() as _,
-            num_bytes: store.metadata_registry.total_size_bytes(),
+        let metadata_registry = {
+            crate::profile_scope!("metadata_registry");
+            DataStoreRowStats {
+                num_rows: store.metadata_registry.len() as _,
+                num_bytes: store.metadata_registry.total_size_bytes(),
+            }
         };
 
-        let autogenerated = DataStoreRowStats {
-            num_rows: store.cluster_cell_cache.len() as _,
-            num_bytes: store.cluster_cell_cache.total_size_bytes(),
+        let autogenerated = {
+            crate::profile_scope!("autogenerated");
+            DataStoreRowStats {
+                num_rows: store.cluster_cell_cache.len() as _,
+                num_bytes: store.cluster_cell_cache.total_size_bytes(),
+            }
         };
 
-        let timeless = DataStoreRowStats {
-            num_rows: store.num_timeless_rows(),
-            num_bytes: store.timeless_size_bytes(),
+        let timeless = {
+            crate::profile_scope!("timeless");
+            DataStoreRowStats {
+                num_rows: store.num_timeless_rows(),
+                num_bytes: store.timeless_size_bytes(),
+            }
         };
 
-        let temporal = DataStoreRowStats {
-            num_rows: store.num_temporal_rows(),
-            num_bytes: store.temporal_size_bytes(),
+        let (temporal, temporal_buckets) = {
+            crate::profile_scope!("temporal");
+            (
+                DataStoreRowStats {
+                    num_rows: store.num_temporal_rows(),
+                    num_bytes: store.temporal_size_bytes(),
+                },
+                store.num_temporal_buckets(),
+            )
         };
-        let temporal_buckets = store.num_temporal_buckets();
 
         let total = DataStoreRowStats {
             num_rows: timeless.num_rows + temporal.num_rows,
@@ -152,7 +169,7 @@ impl SizeBytes for DataTypeRegistry {
 impl SizeBytes for MetadataRegistry<TimePoint> {
     #[inline]
     fn heap_size_bytes(&self) -> u64 {
-        self.0.heap_size_bytes()
+        self.heap_size_bytes
     }
 }
 

crates/re_arrow_store/src/store_write.rs

@@ -233,9 +233,14 @@ impl DataStore {
 
 impl MetadataRegistry<TimePoint> {
     fn upsert(&mut self, row_id: RowId, timepoint: TimePoint) {
+        let mut added_size_bytes = 0;
+
         // This is valuable information even for a timeless timepoint!
         match self.entry(row_id) {
             std::collections::btree_map::Entry::Vacant(entry) => {
+                // NOTE: In a map, thus on the heap!
+                added_size_bytes += row_id.total_size_bytes();
+                added_size_bytes += timepoint.total_size_bytes();
                 entry.insert(timepoint);
             }
             // NOTE: When saving and loading data from disk, it's very possible that we try to
@@ -249,10 +254,16 @@ impl MetadataRegistry<TimePoint> {
                             re_log::error!(%row_id, ?timeline, old_time = ?old_time, new_time = ?time, "detected re-used `RowId/Timeline` pair, this is illegal and will lead to undefined behavior in the datastore");
                             debug_assert!(false, "detected re-used `RowId/Timeline`");
                         }
+                    } else {
+                        // NOTE: In a map, thus on the heap!
+                        added_size_bytes += timeline.total_size_bytes();
+                        added_size_bytes += time.as_i64().total_size_bytes();
                     }
                 }
             }
         }
+
+        self.heap_size_bytes += added_size_bytes;
     }
 }