Merge branch 'main' into benliepert/update_ndarray

benliepert · web-flow · commit 339b4c5558f3 · 2024-09-05T20:53:38.000-04:00
diff --git a/crates/store/re_chunk/src/util.rs b/crates/store/re_chunk/src/util.rs
@@ -1,10 +1,14 @@
 use arrow2::{
-    array::{Array as ArrowArray, BooleanArray as ArrowBooleanArray, ListArray as ArrowListArray},
+    array::{
+        Array as ArrowArray, BooleanArray as ArrowBooleanArray,
+        DictionaryArray as ArrowDictionaryArray, ListArray as ArrowListArray,
+        PrimitiveArray as ArrowPrimitiveArray,
+    },
     bitmap::Bitmap as ArrowBitmap,
-    datatypes::DataType as ArrowDataType,
+    datatypes::DataType as ArrowDatatype,
     offset::Offsets as ArrowOffsets,
 };
-use itertools::Itertools as _;
+use itertools::Itertools;
 
 // ---
 
@@ -29,7 +33,7 @@ pub fn arrays_to_list_array_opt(arrays: &[Option<&dyn ArrowArray>]) -> Option<Ar
 ///
 /// Returns an empty list if `arrays` is empty.
 pub fn arrays_to_list_array(
-    array_datatype: ArrowDataType,
+    array_datatype: ArrowDatatype,
     arrays: &[Option<&dyn ArrowArray>],
 ) -> Option<ArrowListArray<i32>> {
     let arrays_dense = arrays.iter().flatten().copied().collect_vec();
@@ -66,6 +70,91 @@ pub fn arrays_to_list_array(
     ))
 }
 
+/// Create a sparse dictionary-array out of an array of (potentially) duplicated arrays.
+///
+/// The `Idx` is used as primary key to drive the deduplication process.
+/// Returns `None` if any of the specified `arrays` doesn't match the given `array_datatype`.
+///
+/// Returns an empty dictionary if `arrays` is empty.
+//
+// TODO(cmc): Ideally I would prefer to just use the array's underlying pointer as primary key, but
+// this has proved extremely brittle in practice. Maybe once we move to arrow-rs.
+// TODO(cmc): A possible improvement would be to pick the smallest key datatype possible based
+// on the cardinality of the input arrays.
+pub fn arrays_to_dictionary<Idx: Copy + Eq>(
+    array_datatype: ArrowDatatype,
+    arrays: &[Option<(Idx, &dyn ArrowArray)>],
+) -> Option<ArrowDictionaryArray<u32>> {
+    // Dedupe the input arrays based on the given primary key.
+    let arrays_dense_deduped = arrays
+        .iter()
+        .flatten()
+        .copied()
+        .dedup_by(|(lhs_index, _), (rhs_index, _)| lhs_index == rhs_index)
+        .map(|(_index, array)| array)
+        .collect_vec();
+
+    // Compute the keys for the final dictionary, using that same primary key.
+    let keys = {
+        let mut cur_key = 0u32;
+        arrays
+            .iter()
+            .dedup_by_with_count(|lhs, rhs| {
+                lhs.map(|(index, _)| index) == rhs.map(|(index, _)| index)
+            })
+            .flat_map(|(count, value)| {
+                if value.is_some() {
+                    let keys = std::iter::repeat(Some(cur_key)).take(count);
+                    cur_key += 1;
+                    keys
+                } else {
+                    std::iter::repeat(None).take(count)
+                }
+            })
+            .collect_vec()
+    };
+
+    // Concatenate the underlying data as usual, except only the _unique_ values!
+    let data = if arrays_dense_deduped.is_empty() {
+        arrow2::array::new_empty_array(array_datatype.clone())
+    } else {
+        arrow2::compute::concatenate::concatenate(&arrays_dense_deduped)
+            .map_err(|err| {
+                re_log::warn_once!("failed to concatenate arrays: {err}");
+                err
+            })
+            .ok()?
+    };
+
+    // We still need the underlying data to be a list-array, so the dictionary's keys can index
+    // into this list-array.
+    let data = {
+        let datatype = ArrowListArray::<i32>::default_datatype(array_datatype);
+
+        #[allow(clippy::unwrap_used)] // yes, these are indeed lengths
+        let offsets =
+            ArrowOffsets::try_from_lengths(arrays_dense_deduped.iter().map(|array| array.len()))
+                .unwrap();
+
+        ArrowListArray::<i32>::new(datatype, offsets.into(), data, None)
+    };
+
+    let datatype = ArrowDatatype::Dictionary(
+        arrow2::datatypes::IntegerType::UInt32,
+        std::sync::Arc::new(data.data_type().clone()),
+        true, // is_sorted
+    );
+
+    // And finally we build our dictionary, which indexes into our concatenated list-array of
+    // unique values.
+    ArrowDictionaryArray::try_new(
+        datatype,
+        ArrowPrimitiveArray::<u32>::from(keys),
+        data.to_boxed(),
+    )
+    .ok()
+}
+
 /// Given a sparse `ArrowListArray` (i.e. an array with a validity bitmap that contains at least
 /// one falsy value), returns a dense `ArrowListArray` that only contains the non-null values from
 /// the original list.
diff --git a/crates/store/re_chunk_store/src/dataframe.rs b/crates/store/re_chunk_store/src/dataframe.rs
@@ -55,11 +55,11 @@ impl ColumnDescriptor {
     }
 
     #[inline]
-    pub fn to_arrow_field(&self) -> ArrowField {
+    pub fn to_arrow_field(&self, datatype: Option<ArrowDatatype>) -> ArrowField {
         match self {
             Self::Control(descr) => descr.to_arrow_field(),
             Self::Time(descr) => descr.to_arrow_field(),
-            Self::Component(descr) => descr.to_arrow_field(),
+            Self::Component(descr) => descr.to_arrow_field(datatype),
         }
     }
 }
@@ -268,7 +268,7 @@ impl ComponentColumnDescriptor {
     }
 
     #[inline]
-    pub fn to_arrow_field(&self) -> ArrowField {
+    pub fn to_arrow_field(&self, wrapped_datatype: Option<ArrowDatatype>) -> ArrowField {
         let Self {
             entity_path,
             archetype_name,
@@ -278,14 +278,14 @@ impl ComponentColumnDescriptor {
             is_static,
         } = self;
 
+        // NOTE: Only the system doing the actual packing knows the final datatype with all of
+        // its wrappers (is it a component array? is it a list? is it a dict?).
+        let datatype = wrapped_datatype.unwrap_or_else(|| datatype.clone());
+
         // TODO(cmc): figure out who's in charge of adding the outer list layer.
         ArrowField::new(
             component_name.short_name().to_owned(),
-            ArrowDatatype::List(std::sync::Arc::new(ArrowField::new(
-                "item",
-                datatype.clone(),
-                true, /* nullable */
-            ))),
+            datatype,
             false, /* nullable */
         )
         // TODO(#6889): This needs some proper sorbetization -- I just threw these names randomly.
diff --git a/crates/store/re_dataframe/src/latest_at.rs b/crates/store/re_dataframe/src/latest_at.rs
@@ -93,20 +93,10 @@ impl LatestAtQueryHandle<'_> {
     /// [`Self::schema`].
     /// Columns that do not yield any data will still be present in the results, filled with null values.
     pub fn get(&self) -> RecordBatch {
-        re_tracing::profile_function!(format!("{:?}", self.query));
+        re_tracing::profile_function!(format!("{}", self.query));
 
         let columns = self.schema();
 
-        let schema = ArrowSchema {
-            fields: columns
-                .iter()
-                .map(ColumnDescriptor::to_arrow_field)
-                .collect(),
-
-            // TODO(#6889): properly some sorbet stuff we want to get in there at some point.
-            metadata: Default::default(),
-        };
-
         let all_units: HashMap<&ComponentColumnDescriptor, UnitChunkShared> = {
             re_tracing::profile_scope!("queries");
 
@@ -164,6 +154,10 @@ impl LatestAtQueryHandle<'_> {
             }
         }
 
+        // If the query didn't return anything at all, we just want a properly empty Recordbatch with
+        // the right schema.
+        let null_array_length = max_time_per_timeline.get(&self.query.timeline).is_some() as usize;
+
         // NOTE: Keep in mind this must match the ordering specified by `Self::schema`.
         let packed_arrays = {
             re_tracing::profile_scope!("packing");
@@ -186,7 +180,12 @@ impl LatestAtQueryHandle<'_> {
                             .and_then(|(_, chunk)| chunk.timelines().get(&descr.timeline).cloned());
 
                         Some(time_column.map_or_else(
-                            || arrow2::array::new_null_array(descr.datatype.clone(), 1),
+                            || {
+                                arrow2::array::new_null_array(
+                                    descr.datatype.clone(),
+                                    null_array_length,
+                                )
+                            },
                             |time_column| time_column.times_array().to_boxed(),
                         ))
                     }
@@ -196,16 +195,28 @@ impl LatestAtQueryHandle<'_> {
                             .get(descr)
                             .and_then(|chunk| chunk.components().get(&descr.component_name))
                             .map_or_else(
-                                || arrow2::array::new_null_array(descr.datatype.clone(), 1),
+                                || {
+                                    arrow2::array::new_null_array(
+                                        descr.datatype.clone(),
+                                        null_array_length,
+                                    )
+                                },
                                 |list_array| list_array.to_boxed(),
                             ),
                     ),
                 })
-                .collect()
+                .collect_vec()
         };
 
         RecordBatch {
-            schema,
+            schema: ArrowSchema {
+                fields: columns
+                    .iter()
+                    .zip(packed_arrays.iter())
+                    .map(|(descr, arr)| descr.to_arrow_field(Some(arr.data_type().clone())))
+                    .collect(),
+                metadata: Default::default(),
+            },
             data: ArrowChunk::new(packed_arrays),
         }
     }
@@ -225,3 +236,141 @@ impl<'a> LatestAtQueryHandle<'a> {
         })
     }
 }
+
+// ---
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use re_chunk::{ArrowArray, Chunk, EntityPath, RowId, TimeInt, TimePoint, Timeline};
+    use re_chunk_store::{
+        ChunkStore, ChunkStoreConfig, ColumnDescriptor, ComponentColumnDescriptor,
+        LatestAtQueryExpression, TimeColumnDescriptor,
+    };
+    use re_log_types::{example_components::MyPoint, StoreId, StoreKind};
+    use re_query::Caches;
+    use re_types::{
+        components::{Color, Position3D, Radius},
+        Loggable,
+    };
+
+    use crate::QueryEngine;
+
+    #[test]
+    fn empty_yields_empty() {
+        let store = ChunkStore::new(
+            StoreId::random(StoreKind::Recording),
+            ChunkStoreConfig::default(),
+        );
+        let cache = Caches::new(&store);
+        let engine = QueryEngine {
+            store: &store,
+            cache: &cache,
+        };
+
+        let query = LatestAtQueryExpression {
+            entity_path_expr: "/**".into(),
+            timeline: Timeline::log_time(),
+            at: TimeInt::MAX,
+        };
+
+        let entity_path: EntityPath = "/points".into();
+        let columns = vec![
+            ColumnDescriptor::Time(TimeColumnDescriptor {
+                timeline: Timeline::log_time(),
+                datatype: Timeline::log_time().datatype(),
+            }),
+            ColumnDescriptor::Time(TimeColumnDescriptor {
+                timeline: Timeline::log_tick(),
+                datatype: Timeline::log_tick().datatype(),
+            }),
+            ColumnDescriptor::Component(ComponentColumnDescriptor::new::<Position3D>(
+                entity_path.clone(),
+            )),
+            ColumnDescriptor::Component(ComponentColumnDescriptor::new::<Radius>(
+                entity_path.clone(),
+            )),
+            ColumnDescriptor::Component(ComponentColumnDescriptor::new::<Color>(entity_path)),
+        ];
+
+        let handle = engine.latest_at(&query, Some(columns.clone()));
+        let batch = handle.get();
+
+        // The output should be an empty recordbatch with the right schema and empty arrays.
+        assert_eq!(0, batch.num_rows());
+        assert!(itertools::izip!(columns.iter(), batch.schema.fields.iter())
+            .all(|(descr, field)| descr.to_arrow_field(None) == *field));
+        assert!(itertools::izip!(columns.iter(), batch.data.iter())
+            .all(|(descr, array)| descr.datatype() == array.data_type()));
+    }
+
+    #[test]
+    fn static_does_yield() {
+        let mut store = ChunkStore::new(
+            StoreId::random(StoreKind::Recording),
+            ChunkStoreConfig::default(),
+        );
+
+        let entity_path: EntityPath = "/points".into();
+        let chunk = Arc::new(
+            Chunk::builder(entity_path.clone())
+                .with_component_batches(
+                    RowId::new(),
+                    TimePoint::default(),
+                    [&[MyPoint::new(1.0, 1.0), MyPoint::new(2.0, 2.0)] as _],
+                )
+                .build()
+                .unwrap(),
+        );
+        _ = store.insert_chunk(&chunk);
+
+        eprintln!("{store}");
+
+        let cache = Caches::new(&store);
+        let engine = QueryEngine {
+            store: &store,
+            cache: &cache,
+        };
+
+        let query = LatestAtQueryExpression {
+            entity_path_expr: "/**".into(),
+            timeline: Timeline::log_time(),
+            at: TimeInt::MAX,
+        };
+
+        let columns = vec![
+            ColumnDescriptor::Time(TimeColumnDescriptor {
+                timeline: Timeline::log_time(),
+                datatype: Timeline::log_time().datatype(),
+            }),
+            ColumnDescriptor::Time(TimeColumnDescriptor {
+                timeline: Timeline::log_tick(),
+                datatype: Timeline::log_tick().datatype(),
+            }),
+            ColumnDescriptor::Component(ComponentColumnDescriptor::new::<MyPoint>(
+                entity_path.clone(),
+            )),
+            ColumnDescriptor::Component(ComponentColumnDescriptor::new::<Radius>(
+                entity_path.clone(),
+            )),
+            ColumnDescriptor::Component(ComponentColumnDescriptor::new::<Color>(entity_path)),
+        ];
+
+        let handle = engine.latest_at(&query, Some(columns.clone()));
+        let batch = handle.get();
+
+        assert_eq!(1, batch.num_rows());
+        assert_eq!(
+            chunk.components().get(&MyPoint::name()).unwrap().to_boxed(),
+            itertools::izip!(batch.schema.fields.iter(), batch.data.iter())
+                .find_map(
+                    |(field, array)| (field.name == MyPoint::name().short_name())
+                        .then_some(array.clone())
+                )
+                .unwrap()
+        );
+        assert!(itertools::izip!(columns.iter(), batch.schema.fields.iter())
+            .all(|(descr, field)| descr.to_arrow_field(None) == *field));
+    }
+}
diff --git a/crates/store/re_dataframe/src/range.rs b/crates/store/re_dataframe/src/range.rs
