SDK batching/revamp 2.2: homegrown arrow size estimation routines

crates/re_arrow_store/benches/data_store.rs

@@ -269,8 +269,7 @@ fn gc(c: &mut Criterion) {
         (NUM_INSTANCES * NUM_ROWS) as _,
     ));
 
-    let mut table = build_table(NUM_INSTANCES as usize, false);
-    table.compute_all_size_bytes();
+    let table = build_table(NUM_INSTANCES as usize, false);
 
     // Default config
     group.bench_function("default", |b| {
@@ -318,16 +317,19 @@ fn build_table(n: usize, packed: bool) -> DataTable {
             )
         }),
     );
-    if cfg!(debug_assertions) {
-        table.compute_all_size_bytes();
-    }
 
     // Do a serialization roundtrip to pack everything in contiguous memory.
     if packed {
         let (schema, columns) = table.serialize().unwrap();
         table = DataTable::deserialize(TableId::ZERO, &schema, &columns).unwrap();
     }
 
+    // NOTE: Using unsized cells will crash in debug mode, and benchmarks are run for 1 iteration,
+    // in debug mode, by the standard test harness.
+    if cfg!(debug_assertions) {
+        table.compute_all_size_bytes();
+    }
+
     table
 }
 

crates/re_arrow_store/src/test_util.rs

@@ -6,26 +6,22 @@ use crate::{DataStore, DataStoreConfig};
 #[macro_export]
 macro_rules! test_row {
     ($entity:ident @ $frames:tt => $n:expr; [$c0:expr $(,)*]) => {{
-        let mut row = ::re_log_types::DataRow::from_cells1(
+        ::re_log_types::DataRow::from_cells1_sized(
             ::re_log_types::RowId::random(),
             $entity.clone(),
             $frames,
             $n,
             $c0,
-        );
-        row.compute_all_size_bytes();
-        row
+        )
     }};
     ($entity:ident @ $frames:tt => $n:expr; [$c0:expr, $c1:expr $(,)*]) => {{
-        let mut row = ::re_log_types::DataRow::from_cells2(
+        ::re_log_types::DataRow::from_cells2_sized(
             ::re_log_types::RowId::random(),
             $entity.clone(),
             $frames,
             $n,
             ($c0, $c1),
-        );
-        row.compute_all_size_bytes();
-        row
+        )
     }};
 }
 

crates/re_arrow_store/tests/internals.rs

@@ -52,24 +52,22 @@ fn pathological_bucket_topology() {
 
         let timepoint = TimePoint::from([build_frame_nr(frame_nr.into())]);
         for _ in 0..num {
-            let mut row = DataRow::from_cells1(
+            let row = DataRow::from_cells1_sized(
                 RowId::random(),
                 ent_path.clone(),
                 timepoint.clone(),
                 num_instances,
                 build_some_instances(num_instances as _),
             );
-            row.compute_all_size_bytes();
             store_forward.insert_row(&row).unwrap();
 
-            let mut row = DataRow::from_cells1(
+            let row = DataRow::from_cells1_sized(
                 RowId::random(),
                 ent_path.clone(),
                 timepoint.clone(),
                 num_instances,
                 build_some_instances(num_instances as _),
             );
-            row.compute_all_size_bytes();
             store_backward.insert_row(&row).unwrap();
         }
     }
@@ -85,15 +83,13 @@ fn pathological_bucket_topology() {
         let rows = range
             .map(|frame_nr| {
                 let timepoint = TimePoint::from([build_frame_nr(frame_nr.into())]);
-                let mut row = DataRow::from_cells1(
+                DataRow::from_cells1_sized(
                     RowId::random(),
                     ent_path.clone(),
                     timepoint,
                     num_instances,
                     build_some_instances(num_instances as _),
-                );
-                row.compute_all_size_bytes();
-                row
+                )
             })
             .collect::<Vec<_>>();
 

crates/re_log_types/src/data_cell.rs

@@ -546,6 +546,8 @@ impl SizeBytes for DataCell {
         (self.inner.size_bytes > 0)
             .then_some(self.inner.size_bytes)
             .unwrap_or_else(|| {
+                // NOTE: Relying on unsized cells is always a mistake, but it isn't worth crashing
+                // the viewer when in release mode.
                 debug_assert!(
                     false,
                     "called `DataCell::heap_size_bytes() without computing it first"
@@ -591,6 +593,7 @@ fn data_cell_sizes() {
     use arrow2::array::UInt64Array;
 
     // not computed
+    // NOTE: Unsized cells are illegal in debug mode and will flat out crash.
     if !cfg!(debug_assertions) {
         let cell = DataCell::from_arrow(InstanceKey::name(), UInt64Array::from_vec(vec![]).boxed());
         assert_eq!(0, cell.heap_size_bytes());

crates/re_log_types/src/data_row.rs

@@ -316,6 +316,24 @@ impl DataRow {
         Self::try_from_cells(row_id, timepoint, entity_path, num_instances, cells).unwrap()
     }
 
+    /// A helper that combines [`Self::from_cells`] followed by [`Self::compute_all_size_bytes`].
+    ///
+    /// See respective documentations for more information.
+    ///
+    /// Beware: this is costly!
+    pub fn from_cells_sized(
+        row_id: RowId,
+        timepoint: impl Into<TimePoint>,
+        entity_path: impl Into<EntityPath>,
+        num_instances: u32,
+        cells: impl IntoIterator<Item = DataCell>,
+    ) -> Self {
+        let mut this =
+            Self::try_from_cells(row_id, timepoint, entity_path, num_instances, cells).unwrap();
+        this.compute_all_size_bytes();
+        this
+    }
+
     /// Turns the `DataRow` into a single-row [`DataTable`].
     #[inline]
     pub fn into_table(self) -> DataTable {
@@ -428,6 +446,32 @@ impl DataRow {
         )
     }
 
+    /// A helper that combines [`Self::from_cells1`] followed by [`Self::compute_all_size_bytes`].
+    ///
+    /// See respective documentations for more information.
+    ///
+    /// Beware: this is costly!
+    pub fn from_cells1_sized<C0>(
+        row_id: RowId,
+        entity_path: impl Into<EntityPath>,
+        timepoint: impl Into<TimePoint>,
+        num_instances: u32,
+        into_cells: C0,
+    ) -> DataRow
+    where
+        C0: Into<DataCell>,
+    {
+        let mut this = Self::from_cells(
+            row_id,
+            timepoint.into(),
+            entity_path.into(),
+            num_instances,
+            [into_cells.into()],
+        );
+        this.compute_all_size_bytes();
+        this
+    }
+
     pub fn try_from_cells1<C0>(
         row_id: RowId,
         entity_path: impl Into<EntityPath>,
@@ -471,6 +515,36 @@ impl DataRow {
         )
     }
 
+    /// A helper that combines [`Self::from_cells2`] followed by [`Self::compute_all_size_bytes`].
+    ///
+    /// See respective documentations for more information.
+    ///
+    /// Beware: this is costly!
+    pub fn from_cells2_sized<C0, C1>(
+        row_id: RowId,
+        entity_path: impl Into<EntityPath>,
+        timepoint: impl Into<TimePoint>,
+        num_instances: u32,
+        into_cells: (C0, C1),
+    ) -> DataRow
+    where
+        C0: Into<DataCell>,
+        C1: Into<DataCell>,
+    {
+        let mut this = Self::from_cells(
+            row_id,
+            timepoint.into(),
+            entity_path.into(),
+            num_instances,
+            [
+                into_cells.0.into(), //
+                into_cells.1.into(), //
+            ],
+        );
+        this.compute_all_size_bytes();
+        this
+    }
+
     pub fn try_from_cells2<C0, C1>(
         row_id: RowId,
         entity_path: impl Into<EntityPath>,
@@ -521,6 +595,38 @@ impl DataRow {
         )
     }
 
+    /// A helper that combines [`Self::from_cells3`] followed by [`Self::compute_all_size_bytes`].
+    ///
+    /// See respective documentations for more information.
+    ///
+    /// Beware: this is costly!
+    pub fn from_cells3_sized<C0, C1, C2>(
+        row_id: RowId,
+        entity_path: impl Into<EntityPath>,
+        timepoint: impl Into<TimePoint>,
+        num_instances: u32,
+        into_cells: (C0, C1, C2),
+    ) -> DataRow
+    where
+        C0: Into<DataCell>,
+        C1: Into<DataCell>,
+        C2: Into<DataCell>,
+    {
+        let mut this = Self::from_cells(
+            row_id,
+            timepoint.into(),
+            entity_path.into(),
+            num_instances,
+            [
+                into_cells.0.into(), //
+                into_cells.1.into(), //
+                into_cells.2.into(), //
+            ],
+        );
+        this.compute_all_size_bytes();
+        this
+    }
+
     pub fn try_from_cells3<C0, C1, C2>(
         row_id: RowId,
         entity_path: impl Into<EntityPath>,

crates/re_log_types/src/data_table.rs

@@ -1317,7 +1317,7 @@ fn data_table_sizes_unions() {
         );
         assert_eq!(num_bytes, table.heap_size_bytes());
 
-        let err_margin = (num_bytes as f64 * 0.05) as u64;
+        let err_margin = (num_bytes as f64 * 0.01) as u64;
         let num_bytes_min = num_bytes;
         let num_bytes_max = num_bytes + err_margin;
 
@@ -1359,7 +1359,7 @@ fn data_table_sizes_unions() {
             .as_slice(),
         ),
         10_000,     // num_rows
-        27_250_064, // expected_num_bytes
+        27_230_064, // expected_num_bytes
     );
 
     // dense union (varying)

crates/re_log_types/src/data_table_batcher.rs

@@ -512,8 +512,7 @@ mod tests {
         let batcher = DataTableBatcher::new(DataTableBatcherConfig::NEVER).unwrap();
         let tables = batcher.tables();
 
-        let mut table = create_table();
-        table.compute_all_size_bytes();
+        let table = create_table();
         let rows = table.to_rows().collect_vec();
 
         for _ in 0..3 {
@@ -551,9 +550,7 @@ mod tests {
 
     #[test]
     fn num_bytes_trigger() {
-        let mut table = create_table();
-        table.compute_all_size_bytes();
-
+        let table = create_table();
         let rows = table.to_rows().collect_vec();
         let flush_duration = std::time::Duration::from_millis(50);
         let flush_num_bytes = rows
@@ -615,9 +612,7 @@ mod tests {
 
     #[test]
     fn num_rows_trigger() {
-        let mut table = create_table();
-        table.compute_all_size_bytes();
-
+        let table = create_table();
         let rows = table.to_rows().collect_vec();
         let flush_duration = std::time::Duration::from_millis(50);
         let flush_num_rows = rows.len() as u64 - 1;
@@ -675,8 +670,7 @@ mod tests {
 
     #[test]
     fn duration_trigger() {
-        let mut table = create_table();
-        table.compute_all_size_bytes();
+        let table = create_table();
         let rows = table.to_rows().collect_vec();
 
         let flush_duration = Duration::from_millis(50);
@@ -791,6 +785,8 @@ mod tests {
             )
         };
 
-        DataTable::from_rows(TableId::ZERO, [row0, row1, row2])
+        let mut table = DataTable::from_rows(TableId::ZERO, [row0, row1, row2]);
+        table.compute_all_size_bytes();
+        table
     }
 }

crates/re_log_types/src/size_bytes.rs

@@ -1,6 +1,7 @@
 use std::collections::{BTreeMap, HashMap};
 
 use arrow2::datatypes::{DataType, Field};
+use nohash_hasher::IntSet;
 use smallvec::SmallVec;
 
 // ---
@@ -291,6 +292,7 @@ fn estimated_bytes_size(array: &dyn Array) -> usize {
         LargeBinary => dyn_binary!(array, BinaryArray<i64>, i64),
         Utf8 => dyn_binary!(array, Utf8Array<i32>, i32),
         LargeUtf8 => dyn_binary!(array, Utf8Array<i64>, i64),
+        // NOTE: Diverges from upstream.
         List | LargeList => {
             let array = array.as_any().downcast_ref::<ListArray<i32>>().unwrap();
 
@@ -320,6 +322,7 @@ fn estimated_bytes_size(array: &dyn Array) -> usize {
                 .sum::<usize>()
                 + validity_size(array.validity())
         }
+        // NOTE: Diverges from upstream.
         Union => {
             let array = array.as_any().downcast_ref::<UnionArray>().unwrap();
 
@@ -349,14 +352,25 @@ fn estimated_bytes_size(array: &dyn Array) -> usize {
                 }
 
                 let fields = array.fields();
-                let values_start = offsets[0] as usize;
-                let values_end = offsets[offsets.len() - 1] as usize;
-                array
-                    .types()
-                    .iter()
-                    .copied()
-                    .map(|ty| {
-                        fields.get(ty as usize).map_or(0, |x| {
+                let types: IntSet<_> = array.types().iter().copied().collect();
+                types
+                    .into_iter()
+                    .map(|cur_ty| {
+                        let mut indices = array
+                            .types()
+                            .iter()
+                            .enumerate()
+                            .filter_map(|(idx, &ty)| (ty == cur_ty).then_some(idx));
+
+                        let idx_start = indices.next().unwrap_or_default();
+                        let mut idx_end = idx_start;
+                        for idx in indices {
+                            idx_end = idx;
+                        }
+
+                        let values_start = offsets[idx_start] as usize;
+                        let values_end = offsets[idx_end] as usize;
+                        fields.get(cur_ty as usize).map_or(0, |x| {
                             estimated_bytes_size(
                                 x.slice(values_start, values_end - values_start).as_ref(),
                             )

crates/re_query/Cargo.toml

@@ -31,7 +31,7 @@ re_format.workspace = true
 re_log_types = { workspace = true, features = ["arrow_datagen"] }
 
 # External dependencies:
-arrow2 = { workspace = true, features = ["compute_concatenate"] }
+arrow2.workspace = true
 document-features = "0.2"
 itertools = { workspace = true }
 thiserror.workspace = true