diff --git a/crates/re_log_types/src/component_types/instance_key.rs b/crates/re_log_types/src/component_types/instance_key.rs index 4cddbad56db7..30d6b6b7893a 100644 --- a/crates/re_log_types/src/component_types/instance_key.rs +++ b/crates/re_log_types/src/component_types/instance_key.rs @@ -36,6 +36,12 @@ impl InstanceKey { /// for example all points in a point cloud entity. pub const SPLAT: Self = Self(u64::MAX); + #[allow(clippy::should_implement_trait)] + #[inline] + pub fn from_iter(it: impl IntoIterator>) -> Vec { + it.into_iter().map(Into::into).collect::>() + } + /// Are we referring to all instances of the entity (e.g. all points in a point cloud entity)? /// /// The opposite of [`Self::is_specific`]. diff --git a/crates/re_log_types/src/data_cell.rs b/crates/re_log_types/src/data_cell.rs index 568df8fe8928..d8b19a9eabc4 100644 --- a/crates/re_log_types/src/data_cell.rs +++ b/crates/re_log_types/src/data_cell.rs @@ -371,6 +371,36 @@ impl DataCell { { self.try_to_native().unwrap() } + + /// Returns the contents of the cell as an iterator of native optional components. + /// + /// Fails if the underlying arrow data cannot be deserialized into `C`. + // + // TODO(#1694): There shouldn't need to be HRTBs (Higher-Rank Trait Bounds) here. + #[inline] + pub fn try_to_native_opt( + &self, + ) -> DataCellResult> + '_> + where + for<'a> &'a C::ArrayType: IntoIterator, + { + use arrow2_convert::deserialize::arrow_array_deserialize_iterator; + arrow_array_deserialize_iterator(&*self.inner.values).map_err(Into::into) + } + + /// Returns the contents of the cell as an iterator of native optional components. + /// + /// Panics if the underlying arrow data cannot be deserialized into `C`. + /// See [`Self::try_to_native_opt`] for a fallible alternative. + // + // TODO(#1694): There shouldn't need to be HRTBs here. + #[inline] + pub fn to_native_opt(&self) -> impl Iterator> + '_ + where + for<'a> &'a C::ArrayType: IntoIterator, + { + self.try_to_native_opt().unwrap() + } } impl DataCell { diff --git a/crates/re_query/src/dataframe_util.rs b/crates/re_query/src/dataframe_util.rs index c71d48485eb5..2dfc18ba224c 100644 --- a/crates/re_query/src/dataframe_util.rs +++ b/crates/re_query/src/dataframe_util.rs @@ -135,9 +135,9 @@ impl ComponentWithInstances { where for<'a> &'a C0::ArrayType: IntoIterator, { - if C0::name() != self.name { + if C0::name() != self.name() { return Err(QueryError::TypeMismatch { - actual: self.name, + actual: self.name(), requested: C0::name(), }); } @@ -145,8 +145,7 @@ impl ComponentWithInstances { let instance_keys: Vec> = self.iter_instance_keys()?.map(Some).collect_vec(); - let values = - arrow_array_deserialize_iterator::>(self.values.as_ref())?.collect_vec(); + let values = self.values.try_to_native_opt()?.collect_vec(); df_builder2::(&instance_keys, &values) } @@ -160,8 +159,7 @@ where pub fn as_df1(&self) -> crate::Result { let instance_keys = self.primary.iter_instance_keys()?.map(Some).collect_vec(); - let primary_values = - arrow_array_deserialize_iterator(self.primary.values.as_ref())?.collect_vec(); + let primary_values = self.primary.values.try_to_native_opt()?.collect_vec(); df_builder2::(&instance_keys, &primary_values) } @@ -173,8 +171,7 @@ where { let instance_keys = self.primary.iter_instance_keys()?.map(Some).collect_vec(); - let primary_values = - arrow_array_deserialize_iterator(self.primary.values.as_ref())?.collect_vec(); + let primary_values = self.primary.values.try_to_native_opt()?.collect_vec(); let c1_values = self.iter_component::()?.collect_vec(); diff --git a/crates/re_query/src/entity_view.rs b/crates/re_query/src/entity_view.rs index ea3a8c18a1fb..a9bf71e39f1a 100644 --- a/crates/re_query/src/entity_view.rs +++ b/crates/re_query/src/entity_view.rs @@ -1,13 +1,13 @@ use std::{collections::BTreeMap, marker::PhantomData}; -use arrow2::array::{Array, MutableArray, PrimitiveArray}; +use arrow2::array::{Array, PrimitiveArray}; use re_format::arrow; use re_log_types::{ component_types::InstanceKey, external::arrow2_convert::{ deserialize::arrow_array_deserialize_iterator, field::ArrowField, serialize::ArrowSerialize, }, - Component, ComponentName, DeserializableComponent, RowId, SerializableComponent, + Component, ComponentName, DataCell, DeserializableComponent, RowId, SerializableComponent, }; use crate::QueryError; @@ -20,55 +20,54 @@ use crate::QueryError; /// See: [`crate::get_component_with_instances`] #[derive(Clone, Debug)] pub struct ComponentWithInstances { - pub(crate) name: ComponentName, - // TODO(jleibs): Remove optional once the store guarantees this will always exist - pub(crate) instance_keys: Option>, - pub(crate) values: Box, + pub(crate) instance_keys: DataCell, + pub(crate) values: DataCell, } impl ComponentWithInstances { + #[inline] pub fn name(&self) -> ComponentName { - self.name + self.values.component_name() } /// Number of values. 1 for splats. + #[inline] pub fn len(&self) -> usize { - self.values.len() + self.values.num_instances() as _ } + #[inline] pub fn is_empty(&self) -> bool { - self.values.len() == 0 + self.values.is_empty() } /// Iterate over the instance keys /// /// If the instance keys don't exist, generate them based on array-index position of the values + #[inline] pub fn iter_instance_keys(&self) -> crate::Result + '_> { - if let Some(keys) = &self.instance_keys { - let iter = arrow_array_deserialize_iterator::(keys.as_ref())?; - Ok(itertools::Either::Left(iter)) - } else { - let auto_num = (0..self.len()).map(|i| InstanceKey(i as u64)); - Ok(itertools::Either::Right(auto_num)) - } + self.instance_keys + .try_to_native::() + .map_err(Into::into) } /// Iterate over the values and convert them to a native `Component` + #[inline] pub fn iter_values( &self, ) -> crate::Result> + '_> where for<'a> &'a C::ArrayType: IntoIterator, { - if C::name() != self.name { + if C::name() != self.name() { return Err(QueryError::TypeMismatch { - actual: self.name, + actual: self.name(), requested: C::name(), }); } Ok(arrow_array_deserialize_iterator::>( - self.values.as_ref(), + self.values.as_arrow_ref(), )?) } @@ -77,9 +76,9 @@ impl ComponentWithInstances { where for<'a> &'a C::ArrayType: IntoIterator, { - if C::name() != self.name { + if C::name() != self.name() { return Err(QueryError::TypeMismatch { - actual: self.name, + actual: self.name(), requested: C::name(), }); } @@ -96,57 +95,35 @@ impl ComponentWithInstances { /// Look up the value that corresponds to a given `InstanceKey` and return as an arrow `Array` pub fn lookup_arrow(&self, instance_key: &InstanceKey) -> Option> { - let offset = if let Some(instance_keys) = &self.instance_keys { - // If `instance_keys` is set, extract the `PrimitiveArray`, and find - // the index of the value by `binary_search` - - let keys = instance_keys - .as_any() - .downcast_ref::>()? - .values(); - - // If the value is splatted, return the offset of the splat - if keys.len() == 1 && keys[0] == InstanceKey::SPLAT.0 { - 0 - } else { - // Otherwise binary search to find the offset of the instance - keys.binary_search(&instance_key.0).ok()? - } + let keys = self + .instance_keys + .as_arrow_ref() + .as_any() + .downcast_ref::>()? + .values(); + + // If the value is splatted, return the offset of the splat + let offset = if keys.len() == 1 && keys[0] == InstanceKey::SPLAT.0 { + 0 } else { - // If `instance_keys` is not set, then offset is the instance because the implicit - // index is a sequential list - let offset = instance_key.0 as usize; - (offset < self.values.len()).then_some(offset)? + // Otherwise binary search to find the offset of the instance + keys.binary_search(&instance_key.0).ok()? as u32 }; - Some(self.values.slice(offset, 1)) + Some(self.values.as_arrow_ref().slice(offset as _, 1)) } /// Produce a `ComponentWithInstances` from native component types pub fn from_native( - instance_keys: Option<&Vec>, - values: &Vec, - ) -> crate::Result { - use re_log_types::external::arrow2_convert::serialize::arrow_serialize_to_mutable_array; - - let instance_keys = if let Some(keys) = instance_keys { - Some( - arrow_serialize_to_mutable_array::>( - keys, - )? - .as_box(), - ) - } else { - None - }; - - let values = arrow_serialize_to_mutable_array::>(values)?.as_box(); - - Ok(ComponentWithInstances { - name: C::name(), + instance_keys: &[InstanceKey], + values: &[C], + ) -> ComponentWithInstances { + let instance_keys = DataCell::from_native(instance_keys); + let values = DataCell::from_native(values); + ComponentWithInstances { instance_keys, values, - }) + } } } @@ -264,10 +241,10 @@ impl std::fmt::Display for EntityView { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let primary_table = arrow::format_table( [ - self.primary.instance_keys.as_ref().unwrap().as_ref(), - self.primary.values.as_ref(), + self.primary.instance_keys.as_arrow_ref(), + self.primary.values.as_arrow_ref(), ], - ["InstanceId", self.primary.name.as_str()], + ["InstanceId", self.primary.name().as_str()], ); f.write_fmt(format_args!("EntityView:\n{primary_table}")) @@ -294,16 +271,19 @@ where for<'a> &'a Primary::ArrayType: IntoIterator, { /// Iterate over the instance keys + #[inline] pub fn iter_instance_keys(&self) -> crate::Result + '_> { self.primary.iter_instance_keys() } /// Iterate over the primary component values. + #[inline] pub fn iter_primary(&self) -> crate::Result> + '_> { self.primary.iter_values() } /// Iterate over the flattened list of primary component values if any. + #[inline] pub fn iter_primary_flattened(&self) -> impl Iterator + '_ { self.primary .iter_values() @@ -314,6 +294,7 @@ where } /// Check if the entity has a component and its not empty + #[inline] pub fn has_component(&self) -> bool { self.components .get(&C::name()) @@ -337,7 +318,7 @@ where let mut component_instance_key_iter = component.iter_instance_keys()?; let component_value_iter = - arrow_array_deserialize_iterator::>(component.values.as_ref())?; + arrow_array_deserialize_iterator::>(component.values.as_arrow_ref())?; let next_component_instance_key = component_instance_key_iter.next(); @@ -349,51 +330,56 @@ where splatted_component_value: None, })) } else { - let nulls = (0..self.primary.values.len()).map(|_| None); + let nulls = (0..self.primary.values.num_instances()).map(|_| None); Ok(itertools::Either::Right(nulls)) } } /// Helper function to produce an `EntityView` from rust-native `field_types` - pub fn from_native(c0: (Option<&Vec>, &Vec)) -> crate::Result { - let primary = ComponentWithInstances::from_native(c0.0, c0.1)?; - - Ok(Self { + #[inline] + pub fn from_native(c0: (&[InstanceKey], &[Primary])) -> Self { + let primary = ComponentWithInstances::from_native(c0.0, c0.1); + Self { row_id: RowId::ZERO, primary, components: Default::default(), phantom: PhantomData, - }) + } } /// Helper function to produce an `EntityView` from rust-native `field_types` + #[inline] pub fn from_native2( - primary: (Option<&Vec>, &Vec), - component: (Option<&Vec>, &Vec), - ) -> crate::Result + primary: (&[InstanceKey], &[Primary]), + component: (&[InstanceKey], &[C]), + ) -> Self where C: Component + 'static, C: ArrowSerialize + ArrowField, { - let primary = ComponentWithInstances::from_native(primary.0, primary.1)?; - let component_c1 = ComponentWithInstances::from_native(component.0, component.1)?; + let primary = ComponentWithInstances::from_native(primary.0, primary.1); + let component_c1 = ComponentWithInstances::from_native(component.0, component.1); - let components = [(component_c1.name, component_c1)].into(); + let components = [(component_c1.name(), component_c1)].into(); - Ok(Self { + Self { row_id: RowId::ZERO, primary, components, phantom: PhantomData, - }) + } } } #[test] fn lookup_value() { + use arrow2::array::MutableArray; use re_log_types::component_types::{InstanceKey, Point2D, Rect2D}; use re_log_types::external::arrow2_convert::serialize::arrow_serialize_to_mutable_array; - let points = vec![ + + let instance_keys = InstanceKey::from_iter(0..5); + + let points = [ Point2D { x: 1.0, y: 2.0 }, // Point2D { x: 3.0, y: 4.0 }, Point2D { x: 5.0, y: 6.0 }, @@ -401,22 +387,23 @@ fn lookup_value() { Point2D { x: 9.0, y: 10.0 }, ]; - let component = ComponentWithInstances::from_native(None, &points).unwrap(); + let component = + ComponentWithInstances::from_native(instance_keys.as_slice(), points.as_slice()); let missing_value = component.lookup_arrow(&InstanceKey(5)); assert_eq!(missing_value, None); let value = component.lookup_arrow(&InstanceKey(2)).unwrap(); - let expected_point = vec![points[2].clone()]; + let expected_point = [points[2].clone()]; let expected_arrow = - arrow_serialize_to_mutable_array::>(&expected_point) + arrow_serialize_to_mutable_array::(expected_point.as_slice()) .unwrap() .as_box(); assert_eq!(expected_arrow, value); - let instance_keys = vec![ + let instance_keys = [ InstanceKey(17), InstanceKey(47), InstanceKey(48), @@ -424,16 +411,16 @@ fn lookup_value() { InstanceKey(472), ]; - let component = ComponentWithInstances::from_native(Some(&instance_keys), &points).unwrap(); + let component = ComponentWithInstances::from_native(instance_keys.as_slice(), &points); let missing_value = component.lookup_arrow(&InstanceKey(46)); assert_eq!(missing_value, None); let value = component.lookup_arrow(&InstanceKey(99)).unwrap(); - let expected_point = vec![points[3].clone()]; + let expected_point = [points[3].clone()]; let expected_arrow = - arrow_serialize_to_mutable_array::>(&expected_point) + arrow_serialize_to_mutable_array::(expected_point.as_slice()) .unwrap() .as_box(); @@ -460,14 +447,14 @@ fn lookup_value() { #[test] fn lookup_splat() { use re_log_types::component_types::{InstanceKey, Point2D}; - let instances = vec![ + let instances = [ InstanceKey::SPLAT, // ]; - let points = vec![ + let points = [ Point2D { x: 1.0, y: 2.0 }, // ]; - let component = ComponentWithInstances::from_native(Some(&instances), &points).unwrap(); + let component = ComponentWithInstances::from_native(instances.as_slice(), points.as_slice()); // Any instance we look up will return the slatted value let value = component.lookup::(&InstanceKey(1)).unwrap(); diff --git a/crates/re_query/src/lib.rs b/crates/re_query/src/lib.rs index 620e0686e309..e17a237159fc 100644 --- a/crates/re_query/src/lib.rs +++ b/crates/re_query/src/lib.rs @@ -6,8 +6,6 @@ // TODO(jleibs) better crate documentation. -// TODO(cmc): make DataCells first-class citizens in re_query too. - mod entity_view; mod query; mod range; @@ -43,6 +41,9 @@ pub enum QueryError { requested: re_log_types::ComponentName, }, + #[error("Error with one or more the underlying data cells: {0}")] + DataCell(#[from] re_log_types::DataCellError), + #[error("Error converting arrow data")] ArrowError(#[from] arrow2::error::Error), diff --git a/crates/re_query/src/query.rs b/crates/re_query/src/query.rs index e52a1deea003..23ec0c70e353 100644 --- a/crates/re_query/src/query.rs +++ b/crates/re_query/src/query.rs @@ -49,6 +49,8 @@ pub fn get_component_with_instances( ent_path: &EntityPath, component: ComponentName, ) -> crate::Result<(RowId, ComponentWithInstances)> { + debug_assert_eq!(store.cluster_key(), InstanceKey::name()); + let components = [InstanceKey::name(), component]; let (row_id, mut cells) = store @@ -58,12 +60,10 @@ pub fn get_component_with_instances( Ok(( row_id, ComponentWithInstances { - name: component, - instance_keys: cells[0].take().map(|cell| cell.to_arrow()), - values: cells[1] - .take() - .map(|cell| cell.to_arrow()) - .ok_or(QueryError::PrimaryNotFound)?, + // NOTE: The unwrap cannot fail, the cluster key's presence is guaranteed + // by the store. + instance_keys: cells[0].take().unwrap(), + values: cells[1].take().ok_or(QueryError::PrimaryNotFound)?, }, )) } diff --git a/crates/re_query/src/range.rs b/crates/re_query/src/range.rs index 0fbd7925030c..2198d9dedcb2 100644 --- a/crates/re_query/src/range.rs +++ b/crates/re_query/src/range.rs @@ -94,19 +94,19 @@ pub fn range_entity_with_primary<'a, Primary: Component + 'a, const N: usize>( store .range(query, ent_path, components) .map(move |(time, row_id, mut cells)| { - let instance_keys = cells[cluster_col].take().map(|cell| cell.to_arrow()); + // NOTE: The unwrap cannot fail, the cluster key's presence is guaranteed + // by the store. + let instance_keys = cells[cluster_col].take().unwrap(); let is_primary = cells[primary_col].is_some(); let cwis = cells .into_iter() - .enumerate() - .map(|(i, cell)| { + .map(|cell| { cell.map(|cell| { ( row_id, ComponentWithInstances { - name: components[i], instance_keys: instance_keys.clone(), /* shallow */ - values: cell.to_arrow(), + values: cell, }, ) }) diff --git a/crates/re_query/src/visit.rs b/crates/re_query/src/visit.rs index dccdb5050b51..86bdb178a179 100644 --- a/crates/re_query/src/visit.rs +++ b/crates/re_query/src/visit.rs @@ -9,23 +9,24 @@ //! # use re_query::EntityView; //! # use re_log_types::component_types::{ColorRGBA, InstanceKey, Point2D}; //! -//! let points = vec![ +//! let instances = InstanceKey::from_iter(0..3); +//! +//! let points = [ //! Point2D { x: 1.0, y: 2.0 }, //! Point2D { x: 3.0, y: 4.0 }, //! Point2D { x: 5.0, y: 6.0 }, //! ]; //! -//! let colors = vec![ +//! let colors = [ //! ColorRGBA(0), //! ColorRGBA(1), //! ColorRGBA(2), //! ]; //! //! let entity_view = EntityView::from_native2( -//! (None, &points), -//! (None, &colors), -//! ) -//! .unwrap(); +//! (&instances, &points), +//! (&instances, &colors), +//! ); //! //! let mut points_out = Vec::::new(); //! let mut colors_out = Vec::::new(); @@ -38,8 +39,8 @@ //! .ok() //! .unwrap(); //! -//! assert_eq!(points, points_out); -//! assert_eq!(colors, colors_out); +//! assert_eq!(points.as_slice(), points_out.as_slice()); +//! assert_eq!(colors.as_slice(), colors_out.as_slice()); //! ``` use re_log_types::{ diff --git a/crates/re_query/tests/visit_tests.rs b/crates/re_query/tests/visit_tests.rs index 3ad0897e2354..7a401e95a229 100644 --- a/crates/re_query/tests/visit_tests.rs +++ b/crates/re_query/tests/visit_tests.rs @@ -4,12 +4,13 @@ use re_query::{ComponentWithInstances, EntityView}; #[test] fn basic_single_iter() { - let points = vec![ + let instance_keys = InstanceKey::from_iter(0..2); + let points = [ Point2D { x: 1.0, y: 2.0 }, // Point2D { x: 3.0, y: 4.0 }, ]; - let component = ComponentWithInstances::from_native(None, &points).unwrap(); + let component = ComponentWithInstances::from_native(&instance_keys, &points); let results = itertools::izip!( points.into_iter(), @@ -24,25 +25,26 @@ fn basic_single_iter() { #[test] fn implicit_joined_iter() { - let points = vec![ + let instance_keys = InstanceKey::from_iter(0..3); + + let points = [ Point2D { x: 1.0, y: 2.0 }, // Point2D { x: 3.0, y: 4.0 }, Point2D { x: 5.0, y: 6.0 }, ]; - let colors = vec![ + let colors = [ ColorRGBA(0), // ColorRGBA(1), ColorRGBA(2), ]; let entity_view = EntityView::from_native2( - (None, &points), // - (None, &colors), - ) - .unwrap(); + (&instance_keys, &points), // + (&instance_keys, &colors), + ); - let expected_colors = vec![ + let expected_colors = [ Some(ColorRGBA(0)), // Some(ColorRGBA(1)), Some(ColorRGBA(2)), @@ -60,27 +62,28 @@ fn implicit_joined_iter() { #[test] fn implicit_primary_joined_iter() { - let points = vec![ + let point_ids = InstanceKey::from_iter(0..3); + + let points = [ Point2D { x: 1.0, y: 2.0 }, // Point2D { x: 3.0, y: 4.0 }, Point2D { x: 5.0, y: 6.0 }, ]; - let color_ids = vec![ + let color_ids = [ InstanceKey(1), // InstanceKey(2), ]; - let colors = vec![ + let colors = [ ColorRGBA(1), // ColorRGBA(2), ]; let entity_view = EntityView::from_native2( - (None, &points), // - (Some(&color_ids), &colors), - ) - .unwrap(); + (&point_ids, &points), // + (&color_ids, &colors), + ); let expected_colors = vec![ None, // @@ -100,19 +103,21 @@ fn implicit_primary_joined_iter() { #[test] fn implicit_component_joined_iter() { - let point_ids = vec![ + let point_ids = [ InstanceKey(0), // InstanceKey(2), InstanceKey(4), ]; - let points = vec![ + let points = [ Point2D { x: 1.0, y: 2.0 }, // Point2D { x: 3.0, y: 4.0 }, Point2D { x: 5.0, y: 6.0 }, ]; - let colors = vec![ + let color_ids = InstanceKey::from_iter(0..5); + + let colors = [ ColorRGBA(0), // ColorRGBA(1), ColorRGBA(2), @@ -121,10 +126,9 @@ fn implicit_component_joined_iter() { ]; let entity_view = EntityView::from_native2( - (Some(&point_ids), &points), // - (None, &colors), - ) - .unwrap(); + (&point_ids, &points), // + (&color_ids, &colors), + ); let expected_colors = vec![ Some(ColorRGBA(0)), // @@ -175,10 +179,9 @@ fn complex_joined_iter() { ]; let entity_view = EntityView::from_native2( - (Some(&point_ids), &points), // - (Some(&color_ids), &colors), - ) - .unwrap(); + (&point_ids, &points), // + (&color_ids, &colors), + ); let expected_colors = vec![ None, @@ -199,25 +202,19 @@ fn complex_joined_iter() { #[test] fn single_visit() { - let points = vec![ + let instance_keys = InstanceKey::from_iter(0..4); + let points = [ Point2D { x: 1.0, y: 2.0 }, Point2D { x: 3.0, y: 4.0 }, Point2D { x: 5.0, y: 6.0 }, Point2D { x: 7.0, y: 8.0 }, ]; - let entity_view = EntityView::from_native((None, &points)).unwrap(); + let entity_view = EntityView::from_native((&instance_keys, &points)); let mut instance_key_out = Vec::::new(); let mut points_out = Vec::::new(); - let expected_instance = vec![ - InstanceKey(0), // - InstanceKey(1), - InstanceKey(2), - InstanceKey(3), - ]; - entity_view .visit1(|instance_key: InstanceKey, point: Point2D| { instance_key_out.push(instance_key); @@ -226,8 +223,8 @@ fn single_visit() { .ok() .unwrap(); - assert_eq!(instance_key_out, expected_instance); - assert_eq!(points, points_out); + assert_eq!(instance_key_out, instance_keys); + assert_eq!(points.as_slice(), points_out.as_slice()); } #[test] @@ -240,6 +237,8 @@ fn joint_visit() { Point2D { x: 9.0, y: 10.0 }, ]; + let point_ids = InstanceKey::from_iter(0..5); + let colors = vec![ ColorRGBA(0xff000000), // ColorRGBA(0x00ff0000), @@ -251,10 +250,9 @@ fn joint_visit() { ]; let entity_view = EntityView::from_native2( - (None, &points), // - (Some(&color_ids), &colors), - ) - .unwrap(); + (&point_ids, &points), // + (&color_ids, &colors), + ); let mut points_out = Vec::::new(); let mut colors_out = Vec::>::new();