From 00eeacbf5a145c12e50122e6152f7e36b039487b Mon Sep 17 00:00:00 2001 From: Jeremy Leibs Date: Thu, 20 Jul 2023 19:59:40 -0400 Subject: [PATCH] Introduce support for querying Archetypes (#2743) ### What The main change here is to introduce a new `ArchetypeView` that uses the archetype interfaces as an alternative to the historical `EntityView`. Because EntityViews and legacy type definitions leak through numerous APIs, this involved forking a number of APIs: - `process_archetype_views` - `query_archetype_with_history` - `query_archetype` - `range_archetype` - `process_color_arch` - `process_radii_arch` - `process_annotations_and_keypoints_arch` As we migrate more of our queries to the new APIs, the legacy versions of these helpers can be incrementally deleted and potentially renamed as appropriate. ### Future work This does not yet replicated the `query_tests` and `range_tests` for Archetype APIs. Part of the reason for this is that we make heavy use of `DataCell` construction in those tests, but DataCells are still heavily dependent on the legacy Component APIs. ### Checklist * [x] I have read and agree to [Contributor Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and the [Code of Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md) * [x] I've included a screenshot or gif (if applicable) * [x] I have tested [demo.rerun.io](https://demo.rerun.io/pr/2743) (if applicable) - [PR Build Summary](https://build.rerun.io/pr/2743) - [Docs preview](https://rerun.io/preview/pr%3Ajleibs%2Fhope_queries/docs) - [Examples preview](https://rerun.io/preview/pr%3Ajleibs%2Fhope_queries/examples) --- Cargo.lock | 5 + crates/re_components/Cargo.toml | 1 + crates/re_components/src/class_id.rs | 6 + crates/re_components/src/keypoint_id.rs | 6 + crates/re_log_types/Cargo.toml | 1 + crates/re_log_types/src/instance_key.rs | 6 + crates/re_query/Cargo.toml | 2 + crates/re_query/src/archetype_view.rs | 442 ++++++++++++++++++ crates/re_query/src/dataframe_util.rs | 98 +++- crates/re_query/src/lib.rs | 26 +- crates/re_query/src/query.rs | 124 ++++- crates/re_query/src/range.rs | 145 +++++- crates/re_query/src/util.rs | 33 +- .../re_query/tests/archetype_visit_tests.rs | 283 +++++++++++ crates/re_space_view_spatial/Cargo.toml | 1 + .../src/parts/entity_iterator.rs | 74 ++- crates/re_space_view_spatial/src/parts/mod.rs | 95 ++++ .../src/parts/points2d.rs | 100 ++-- crates/re_types/source_hash.txt | 2 +- crates/re_types/src/archetypes/fuzzy.rs | 7 + crates/re_types/src/archetypes/points2d.rs | 39 +- crates/re_types/src/archetypes/transform3d.rs | 7 + .../src/components/instance_key_ext.rs | 12 + crates/re_types/src/lib.rs | 5 +- crates/re_types_builder/src/codegen/rust.rs | 29 +- 25 files changed, 1469 insertions(+), 80 deletions(-) create mode 100644 crates/re_query/src/archetype_view.rs create mode 100644 crates/re_query/tests/archetype_visit_tests.rs diff --git a/Cargo.lock b/Cargo.lock index d917b332997b..06ef15ba16f0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3991,6 +3991,7 @@ dependencies = [ "re_log", "re_log_types", "re_tracing", + "re_types", "rmp-serde", "serde", "similar-asserts", @@ -4152,6 +4153,7 @@ dependencies = [ "re_string_interner", "re_tracing", "re_tuid", + "re_types", "rmp-serde", "serde", "serde_bytes", @@ -4189,6 +4191,7 @@ name = "re_query" version = "0.8.0-alpha.6" dependencies = [ "arrow2", + "backtrace", "criterion", "document-features", "itertools 0.11.0", @@ -4201,6 +4204,7 @@ dependencies = [ "re_log", "re_log_types", "re_tracing", + "re_types", "thiserror", ] @@ -4366,6 +4370,7 @@ dependencies = [ "re_renderer", "re_space_view", "re_tracing", + "re_types", "re_ui", "re_viewer_context", "serde", diff --git a/crates/re_components/Cargo.toml b/crates/re_components/Cargo.toml index a55f70ace334..46d4df6d2bd1 100644 --- a/crates/re_components/Cargo.toml +++ b/crates/re_components/Cargo.toml @@ -42,6 +42,7 @@ serde = ["dep:serde", "half/serde", "re_log_types/serde"] re_log_types.workspace = true re_log.workspace = true re_tracing.workspace = true +re_types.workspace = true # External ahash.workspace = true diff --git a/crates/re_components/src/class_id.rs b/crates/re_components/src/class_id.rs index af6f6d98ee96..a049baa0d864 100644 --- a/crates/re_components/src/class_id.rs +++ b/crates/re_components/src/class_id.rs @@ -27,3 +27,9 @@ impl re_log_types::Component for ClassId { "rerun.class_id".into() } } + +impl From for ClassId { + fn from(other: re_types::components::ClassId) -> Self { + Self(other.0) + } +} diff --git a/crates/re_components/src/keypoint_id.rs b/crates/re_components/src/keypoint_id.rs index edc2133ffed9..3f625faa672e 100644 --- a/crates/re_components/src/keypoint_id.rs +++ b/crates/re_components/src/keypoint_id.rs @@ -29,3 +29,9 @@ impl re_log_types::Component for KeypointId { "rerun.keypoint_id".into() } } + +impl From for KeypointId { + fn from(other: re_types::components::KeypointId) -> Self { + Self(other.0) + } +} diff --git a/crates/re_log_types/Cargo.toml b/crates/re_log_types/Cargo.toml index e992b275212d..371166267eb0 100644 --- a/crates/re_log_types/Cargo.toml +++ b/crates/re_log_types/Cargo.toml @@ -36,6 +36,7 @@ re_format.workspace = true re_log.workspace = true re_string_interner.workspace = true re_tracing.workspace = true +re_types.workspace = true re_tuid = { workspace = true, features = ["arrow2_convert"] } # External diff --git a/crates/re_log_types/src/instance_key.rs b/crates/re_log_types/src/instance_key.rs index 6ef75a621f0e..966ad8741c85 100644 --- a/crates/re_log_types/src/instance_key.rs +++ b/crates/re_log_types/src/instance_key.rs @@ -104,3 +104,9 @@ impl Component for InstanceKey { "rerun.instance_key".into() } } + +impl From for InstanceKey { + fn from(other: re_types::components::InstanceKey) -> Self { + Self(other.0) + } +} diff --git a/crates/re_query/Cargo.toml b/crates/re_query/Cargo.toml index c899f242f99d..5262d5bd601c 100644 --- a/crates/re_query/Cargo.toml +++ b/crates/re_query/Cargo.toml @@ -30,11 +30,13 @@ re_components = { workspace = true, features = ["arrow_datagen"] } re_data_store.workspace = true re_format.workspace = true re_log_types.workspace = true +re_types.workspace = true re_log.workspace = true re_tracing.workspace = true # External dependencies: arrow2.workspace = true +backtrace = "0.3" document-features = "0.2" itertools = { workspace = true } thiserror.workspace = true diff --git a/crates/re_query/src/archetype_view.rs b/crates/re_query/src/archetype_view.rs new file mode 100644 index 000000000000..bb88b56b6ab9 --- /dev/null +++ b/crates/re_query/src/archetype_view.rs @@ -0,0 +1,442 @@ +use std::{collections::BTreeMap, marker::PhantomData}; + +use arrow2::array::{Array, PrimitiveArray}; +use re_format::arrow; +use re_log_types::{DataCell, RowId}; +use re_types::{ + components::InstanceKey, Archetype, Component, ComponentName, DeserializationResult, Loggable, +}; + +use crate::{ComponentWithInstances, QueryError}; + +/// A type-erased array of [`Component`] values and the corresponding [`InstanceKey`] keys. +/// +/// See: [`crate::get_component_with_instances`] +#[derive(Clone, Debug)] +pub struct ArchComponentWithInstances { + pub(crate) instance_keys: DataCell, + pub(crate) values: DataCell, +} + +impl From for ArchComponentWithInstances { + fn from(other: ComponentWithInstances) -> Self { + ArchComponentWithInstances { + instance_keys: other.instance_keys, + values: other.values, + } + } +} + +impl ArchComponentWithInstances { + /// Name of the [`Component`] + #[inline] + pub fn name(&self) -> ComponentName { + self.values.component_name().as_str().into() + } + + /// Number of values. 1 for splats. + #[inline] + pub fn len(&self) -> usize { + self.values.num_instances() as _ + } + + #[inline] + /// Whether this [`ArchComponentWithInstances`] contains any data + pub fn is_empty(&self) -> bool { + self.values.is_empty() + } + + /// Iterate over the [`InstanceKey`]s. + #[inline] + pub fn iter_instance_keys(&self) -> impl Iterator + '_ { + InstanceKey::from_arrow(self.instance_keys.as_arrow_ref()).into_iter() + } + + /// Iterate over the values and convert them to a native [`Component`] + #[inline] + pub fn iter_values<'a, C: Component + 'a>( + &self, + ) -> crate::Result> + 'a> { + if C::name() != self.name() { + return Err(QueryError::NewTypeMismatch { + actual: self.name().clone(), + requested: C::name(), + }); + } + + Ok(C::try_from_arrow_opt(self.values.as_arrow_ref())?.into_iter()) + } + + /// Look up the value that corresponds to a given [`InstanceKey`] and convert to [`Component`] + pub fn lookup(&self, instance_key: &InstanceKey) -> crate::Result { + if C::name() != self.name() { + return Err(QueryError::NewTypeMismatch { + actual: self.name().clone(), + requested: C::name(), + }); + } + let arr = self + .lookup_arrow(instance_key) + .map_or_else(|| Err(QueryError::ComponentNotFound), Ok)?; + + let mut iter = C::try_from_arrow(arr.as_ref())?.into_iter(); + + let val = iter + .next() + .map_or_else(|| Err(QueryError::ComponentNotFound), Ok)?; + Ok(val) + } + + /// 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 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 { + // Otherwise binary search to find the offset of the instance + keys.binary_search(&instance_key.0).ok()? + }; + + (self.len() > offset) + .then(|| self.values.as_arrow_ref().sliced(offset, 1)) + .or_else(|| { + re_log::error_once!("found corrupt cell -- mismatched number of instances"); + None + }) + } + + /// Produce a [`ArchComponentWithInstances`] from native [`Component`] types + pub fn from_native<'a, C: Component + Clone + 'a>( + instance_keys: impl IntoIterator>>, + values: impl IntoIterator>>, + ) -> ArchComponentWithInstances { + let instance_keys = InstanceKey::to_arrow(instance_keys, None); + let values = C::to_arrow(values, None); + ArchComponentWithInstances { + instance_keys: DataCell::from_arrow(InstanceKey::name().as_ref().into(), instance_keys), + values: DataCell::from_arrow(C::name().as_ref().into(), values), + } + } +} + +/// Iterator over a single [`Component`] joined onto a primary [`Component`] +/// +/// This is equivalent to a left join between one table made up of the +/// [`InstanceKey`]s from the primary component and another table with the +/// [`InstanceKey`]s and values of the iterated [`Component`]. +/// +/// Instances have a [`InstanceKey::SPLAT`] key that will cause the value to be +/// repeated for the entirety of the join. +/// +/// For example +/// ```text +/// primary +/// +----------+ +/// | instance | +/// +----------+ +/// | key0 | +/// | key1 | +/// | Key2 | +/// +/// component +/// +----------+-------+ +/// | instance | value | +/// +----------+-------+ +/// | key0 | val0 | +/// | Key2 | val2 | +/// +/// SELECT value FROM LEFT JOIN primary.instance = component.instance; +/// +/// output +/// +-------+ +/// | value | +/// +-------+ +/// | val0 | +/// | NULL | +/// | val2 | +/// +/// ``` +struct ArchComponentJoinedIterator { + primary_instance_key_iter: IIter1, + component_instance_key_iter: IIter2, + component_value_iter: VIter, + next_component_instance_key: Option, + splatted_component_value: Option, +} + +impl Iterator for ArchComponentJoinedIterator +where + IIter1: Iterator, + IIter2: Iterator, + VIter: Iterator>, + C: Clone, +{ + type Item = Option; + + fn next(&mut self) -> Option> { + // For each value of primary_instance_iter we must find a result + if let Some(primary_key) = self.primary_instance_key_iter.next() { + loop { + match &self.next_component_instance_key { + // If we have a next component key, we either... + Some(instance_key) => { + if instance_key.is_splat() { + if self.splatted_component_value.is_none() { + self.splatted_component_value = + self.component_value_iter.next().flatten(); + } + break Some(self.splatted_component_value.clone()); + } else { + match primary_key.0.cmp(&instance_key.0) { + // Return a None if the primary_key hasn't reached it yet + std::cmp::Ordering::Less => break Some(None), + // Return the value if the keys match + std::cmp::Ordering::Equal => { + self.next_component_instance_key = + self.component_instance_key_iter.next(); + break self.component_value_iter.next(); + } + // Skip this component if the key is behind the primary key + std::cmp::Ordering::Greater => { + _ = self.component_value_iter.next(); + self.next_component_instance_key = + self.component_instance_key_iter.next(); + } + } + } + } + // Otherwise, we ran out of component elements. Just return + // None until the primary iter ends. + None => break Some(None), + }; + } + } else { + None + } + } +} + +/// A view of an [`Archetype`] at a particular point in time returned by [`crate::get_component_with_instances`] +/// +/// The required [`Component`]s of an [`ArchetypeView`] determines the length of an entity +/// batch. When iterating over individual components, they will be implicitly joined onto +/// the required [`Component`]s using [`InstanceKey`] values. +#[derive(Clone, Debug)] +pub struct ArchetypeView { + pub(crate) row_id: RowId, + pub(crate) components: BTreeMap, + pub(crate) phantom: PhantomData, +} + +impl std::fmt::Display for ArchetypeView { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let first_required = self.required_comp(); + + let primary_table = arrow::format_table( + [ + first_required.instance_keys.as_arrow_ref(), + first_required.values.as_arrow_ref(), + ], + ["InstanceId", first_required.name().as_ref()], + ); + + f.write_fmt(format_args!("ArchetypeView:\n{primary_table}")) + } +} + +impl ArchetypeView { + #[inline] + pub fn num_instances(&self) -> usize { + self.required_comp().len() + } + + #[inline] + pub fn row_id(&self) -> RowId { + self.row_id + } +} + +impl ArchetypeView { + fn required_comp(&self) -> &ArchComponentWithInstances { + // TODO(jleibs): Do all archetypes always have at least 1 required components? + let first_required = A::required_components()[0].clone(); + &self.components[&first_required] + } + + /// Iterate over the [`InstanceKey`]s. + #[inline] + pub fn iter_instance_keys(&self) -> impl Iterator + '_ { + // TODO(https://github.com/rerun-io/rerun/issues/2750): Maybe make this an intersection instead + self.required_comp().iter_instance_keys() + } + + /// Check if the entity has a component and its not empty + #[inline] + pub fn has_component(&self) -> bool { + let name = C::name(); + self.components.get(&name).map_or(false, |c| !c.is_empty()) + } + + /// Iterate over the values of a required [`Component`]. + pub fn iter_required_component<'a, C: Component + Clone + 'a>( + &'a self, + ) -> DeserializationResult + '_> { + debug_assert!(A::required_components() + .iter() + .any(|c| c.as_ref() == C::name())); + let component = self.components.get(&C::name()); + + if let Some(component) = component { + let component_value_iter = + C::try_from_arrow(component.values.as_arrow_ref())?.into_iter(); + + Ok(component_value_iter) + } else { + Err(re_types::DeserializationError::MissingData { + backtrace: ::backtrace::Backtrace::new_unresolved(), + }) + } + } + + /// Iterate over the values of an optional `Component`. + /// + /// Always produces an iterator that matches the length of a primary + /// component by joining on the `InstanceKey` values. + pub fn iter_optional_component<'a, C: Component + Clone + 'a>( + &'a self, + ) -> DeserializationResult> + '_> { + let component = self.components.get(&C::name()); + + if let Some(component) = component { + let primary_instance_key_iter = self.iter_instance_keys(); + + let mut component_instance_key_iter = component.iter_instance_keys(); + + let component_value_iter = + C::try_from_arrow_opt(component.values.as_arrow_ref())?.into_iter(); + + let next_component_instance_key = component_instance_key_iter.next(); + + Ok(itertools::Either::Left(ArchComponentJoinedIterator { + primary_instance_key_iter, + component_instance_key_iter, + component_value_iter, + next_component_instance_key, + splatted_component_value: None, + })) + } else { + let primary = self.required_comp(); + let nulls = (0..primary.len()).map(|_| None); + Ok(itertools::Either::Right(nulls)) + } + } + + /// Helper function to produce an [`ArchetypeView`] from a collection of [`ArchComponentWithInstances`] + #[inline] + pub fn from_components( + components: impl IntoIterator, + ) -> Self { + Self { + row_id: RowId::ZERO, + components: components + .into_iter() + .map(|comp| (comp.name().clone(), comp)) + .collect(), + phantom: PhantomData, + } + } +} + +#[test] +fn lookup_value() { + use re_types::components::{InstanceKey, Point2D}; + + let instance_keys = InstanceKey::from_iter(0..5); + + let points = [ + Point2D::new(1.0, 2.0), // + Point2D::new(3.0, 4.0), + Point2D::new(5.0, 6.0), + Point2D::new(7.0, 8.0), + Point2D::new(9.0, 10.0), + ]; + + let component = ArchComponentWithInstances::from_native(instance_keys, points); + + let missing_value = component.lookup_arrow(&InstanceKey(5)); + assert_eq!(missing_value, None); + + let value = component.lookup_arrow(&InstanceKey(2)).unwrap(); + + let expected_point = [points[2]]; + let expected_arrow = Point2D::to_arrow(expected_point, None); + + assert_eq!(expected_arrow, value); + + let instance_keys = [ + InstanceKey(17), + InstanceKey(47), + InstanceKey(48), + InstanceKey(99), + InstanceKey(472), + ]; + + let component = ArchComponentWithInstances::from_native(instance_keys, 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 = [points[3]]; + let expected_arrow = Point2D::to_arrow(expected_point, None); + + assert_eq!(expected_arrow, value); + + // Lookups with serialization + + let value = component.lookup::(&InstanceKey(99)).unwrap(); + assert_eq!(expected_point[0], value); + + let missing_value = component.lookup::(&InstanceKey(46)); + assert!(matches!( + missing_value.err().unwrap(), + QueryError::ComponentNotFound + )); + + // TODO(jleibs): Add another type + /* + let missing_value = component.lookup::(&InstanceKey(99)); + assert!(matches!( + missing_value.err().unwrap(), + QueryError::TypeMismatch { .. } + )); + */ +} + +#[test] +fn lookup_splat() { + use re_types::components::{InstanceKey, Point2D}; + let instances = [ + InstanceKey::SPLAT, // + ]; + let points = [ + Point2D::new(1.0, 2.0), // + ]; + + let component = ArchComponentWithInstances::from_native(instances, points); + + // Any instance we look up will return the slatted value + let value = component.lookup::(&InstanceKey(1)).unwrap(); + assert_eq!(points[0], value); + + let value = component.lookup::(&InstanceKey(99)).unwrap(); + assert_eq!(points[0], value); +} diff --git a/crates/re_query/src/dataframe_util.rs b/crates/re_query/src/dataframe_util.rs index 492b0621fe86..0e78e8516a0f 100644 --- a/crates/re_query/src/dataframe_util.rs +++ b/crates/re_query/src/dataframe_util.rs @@ -9,10 +9,11 @@ use re_log_types::{ external::arrow2_convert::deserialize::arrow_array_deserialize_iterator, Component, DeserializableComponent, InstanceKey, SerializableComponent, }; +use re_types::{Archetype, Loggable}; use crate::{ entity_view::{ComponentWithInstances, EntityView}, - QueryError, + ArchetypeView, QueryError, }; /// Make it so that our arrays can be deserialized again by arrow2-convert @@ -127,6 +128,61 @@ where Ok(DataFrame::new(vec![series0, series1, series2])?) } +pub fn df_builder1_arch<'a, C0>(c0: &'a [Option]) -> crate::Result +where + C0: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C0>> + 'a, + &'a C0: Into<::std::borrow::Cow<'a, C0>>, +{ + let array0 = C0::try_to_arrow_opt(c0.iter().map(|c| c.as_ref()), None)?; + + let series0 = Series::try_from((C0::name().as_ref(), array0.as_ref().clean_for_polars()))?; + + Ok(DataFrame::new(vec![series0])?) +} + +pub fn df_builder2_arch<'a, C0, C1>( + c0: &'a [Option], + c1: &'a [Option], +) -> crate::Result +where + C0: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C0>> + 'a, + C1: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C1>> + 'a, + &'a C0: Into<::std::borrow::Cow<'a, C0>>, + &'a C1: Into<::std::borrow::Cow<'a, C1>>, +{ + let array0 = C0::try_to_arrow_opt(c0.iter().map(|c| c.as_ref()), None)?; + let array1 = C1::try_to_arrow_opt(c1.iter().map(|c| c.as_ref()), None)?; + + let series0 = Series::try_from((C0::name().as_ref(), array0.as_ref().clean_for_polars()))?; + let series1 = Series::try_from((C1::name().as_ref(), array1.as_ref().clean_for_polars()))?; + + Ok(DataFrame::new(vec![series0, series1])?) +} + +pub fn df_builder3_arch<'a, C0, C1, C2>( + c0: &'a [Option], + c1: &'a [Option], + c2: &'a [Option], +) -> crate::Result +where + C0: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C0>> + 'a, + C1: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C1>> + 'a, + C2: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C2>> + 'a, + &'a C0: Into<::std::borrow::Cow<'a, C0>>, + &'a C1: Into<::std::borrow::Cow<'a, C1>>, + &'a C2: Into<::std::borrow::Cow<'a, C2>>, +{ + let array0 = C0::try_to_arrow_opt(c0.iter().map(|c| c.as_ref()), None)?; + let array1 = C1::try_to_arrow_opt(c1.iter().map(|c| c.as_ref()), None)?; + let array2 = C2::try_to_arrow_opt(c2.iter().map(|c| c.as_ref()), None)?; + + let series0 = Series::try_from((C0::name().as_ref(), array0.as_ref().clean_for_polars()))?; + let series1 = Series::try_from((C1::name().as_ref(), array1.as_ref().clean_for_polars()))?; + let series2 = Series::try_from((C2::name().as_ref(), array2.as_ref().clean_for_polars()))?; + + Ok(DataFrame::new(vec![series0, series1, series2])?) +} + impl ComponentWithInstances { pub fn as_df( &self, @@ -178,6 +234,46 @@ where } } +impl ArchetypeView { + pub fn as_df1<'a, C1: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C1>> + 'a>( + &self, + ) -> crate::Result { + let array0 = + re_types::components::InstanceKey::try_to_arrow(self.iter_instance_keys(), None)?; + let array1 = C1::try_to_arrow_opt(self.iter_optional_component::()?, None)?; + + let series0 = Series::try_from(( + re_types::components::InstanceKey::name().as_ref(), + array0.as_ref().clean_for_polars(), + ))?; + let series1 = Series::try_from((C1::name().as_ref(), array1.as_ref().clean_for_polars()))?; + + Ok(DataFrame::new(vec![series0, series1])?) + } + + pub fn as_df2< + 'a, + C1: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C1>> + 'a, + C2: re_types::Component + Clone + Into<::std::borrow::Cow<'a, C2>> + 'a, + >( + &self, + ) -> crate::Result { + let array0 = + re_types::components::InstanceKey::try_to_arrow(self.iter_instance_keys(), None)?; + let array1 = C1::try_to_arrow_opt(self.iter_optional_component::()?, None)?; + let array2 = C2::try_to_arrow_opt(self.iter_optional_component::()?, None)?; + + let series0 = Series::try_from(( + re_types::components::InstanceKey::name().as_ref(), + array0.as_ref().clean_for_polars(), + ))?; + let series1 = Series::try_from((C1::name().as_ref(), array1.as_ref().clean_for_polars()))?; + let series2 = Series::try_from((C2::name().as_ref(), array2.as_ref().clean_for_polars()))?; + + Ok(DataFrame::new(vec![series0, series1, series2])?) + } +} + #[test] fn test_df_builder() { use re_components::{ColorRGBA, Point2D}; diff --git a/crates/re_query/src/lib.rs b/crates/re_query/src/lib.rs index b32d11d7509e..a593c2d02dfb 100644 --- a/crates/re_query/src/lib.rs +++ b/crates/re_query/src/lib.rs @@ -6,6 +6,7 @@ // TODO(jleibs) better crate documentation. +mod archetype_view; mod entity_view; mod query; mod range; @@ -15,10 +16,11 @@ pub mod visit; #[cfg(feature = "polars")] pub mod dataframe_util; +pub use self::archetype_view::{ArchComponentWithInstances, ArchetypeView}; pub use self::entity_view::{ComponentWithInstances, EntityView}; -pub use self::query::{get_component_with_instances, query_entity_with_primary}; -pub use self::range::range_entity_with_primary; -pub use self::util::query_primary_with_history; +pub use self::query::{get_component_with_instances, query_archetype, query_entity_with_primary}; +pub use self::range::{range_archetype, range_entity_with_primary}; +pub use self::util::{query_archetype_with_history, query_primary_with_history}; // Used for doc-tests #[doc(hidden)] @@ -32,6 +34,9 @@ pub enum QueryError { #[error("Could not find primary")] PrimaryNotFound, + #[error("Could not find required component: {0}")] + RequiredComponentNotFound(re_log_types::ComponentName), + #[error("Could not find component")] ComponentNotFound, @@ -41,15 +46,30 @@ pub enum QueryError { requested: re_log_types::ComponentName, }, + #[error("Tried to access component of type '{actual:?}' using component '{requested:?}'")] + NewTypeMismatch { + actual: re_types::ComponentName, + requested: re_types::ComponentName, + }, + #[error("Error with one or more the underlying data cells: {0}")] DataCell(#[from] re_log_types::DataCellError), + #[error("Error deserializing: {0}")] + DeserializationError(#[from] re_types::DeserializationError), + + #[error("Error serializing: {0}")] + SerializationError(#[from] re_types::SerializationError), + #[error("Error converting arrow data: {0}")] ArrowError(#[from] arrow2::error::Error), #[cfg(feature = "polars")] #[error("Error from within Polars")] PolarsError(#[from] polars_core::prelude::PolarsError), + + #[error("Not implemented")] + NotImplemented, } pub type Result = std::result::Result; diff --git a/crates/re_query/src/query.rs b/crates/re_query/src/query.rs index bfa62d5d4e0e..c07917994d33 100644 --- a/crates/re_query/src/query.rs +++ b/crates/re_query/src/query.rs @@ -2,8 +2,9 @@ use std::collections::BTreeMap; use re_arrow_store::{DataStore, LatestAtQuery}; use re_log_types::{Component, ComponentName, DataRow, EntityPath, InstanceKey, RowId}; +use re_types::Archetype; -use crate::{ComponentWithInstances, EntityView, QueryError}; +use crate::{ArchetypeView, ComponentWithInstances, EntityView, QueryError}; /// Retrieves a [`ComponentWithInstances`] from the [`DataStore`]. /// @@ -152,6 +153,89 @@ pub fn query_entity_with_primary( }) } +/// Retrieve an [`ArchetypeView`] from the `DataStore` +/// +/// If you expect only one instance (e.g. for mono-components like `Transform` `Tensor`] +/// and have no additional components you can use [`DataStore::query_latest_component`] instead. +/// +/// ``` +/// # use re_arrow_store::LatestAtQuery; +/// # use re_log_types::Timeline; +/// # use re_types::Component; +/// # use re_types::components::{Point2D, Color}; +/// # use re_types::archetypes::Points2D; +/// # let store = re_query::__populate_example_store(); +/// +/// let ent_path = "point"; +/// let query = LatestAtQuery::new(Timeline::new_sequence("frame_nr"), 123.into()); +/// +/// let arch_view = re_query::query_archetype::( +/// &store, +/// &query, +/// &ent_path.into(), +/// ) +/// .unwrap(); +/// +/// # #[cfg(feature = "polars")] +/// let df = arch_view.as_df2::().unwrap(); +/// +/// //println!("{df:?}"); +/// ``` +/// +/// Outputs: +/// ```text +/// ┌────────────────────┬───────────────┬─────────────────┐ +/// │ rerun.instance_key ┆ rerun.point2d ┆ rerun.colorrgba │ +/// │ --- ┆ --- ┆ --- │ +/// │ u64 ┆ struct[2] ┆ u32 │ +/// ╞════════════════════╪═══════════════╪═════════════════╡ +/// │ 42 ┆ {1.0,2.0} ┆ null │ +/// │ 96 ┆ {3.0,4.0} ┆ 4278190080 │ +/// └────────────────────┴───────────────┴─────────────────┘ +/// ``` +/// +pub fn query_archetype( + store: &DataStore, + query: &LatestAtQuery, + ent_path: &EntityPath, +) -> crate::Result> { + re_tracing::profile_function!(); + + let required_components: Vec<_> = A::required_components() + .iter() + .map(|component| { + get_component_with_instances(store, query, ent_path, component.as_ref().into()) + .map(|(_, component_result)| component_result.into()) + }) + .collect(); + + // NOTE: It's important to use `PrimaryNotFound` here. Any other error will be + // reported to the user. + // + // `query_archetype` is currently run for every archetype on every path in the view + // each path that's missing the primary is then ignored rather than being visited. + if required_components.iter().any(|c| c.is_none()) { + return crate::Result::Err(QueryError::PrimaryNotFound); + } + + let required_components = required_components.into_iter().flatten(); + + let recommended_components = A::recommended_components(); + let optional_components = A::optional_components(); + + let other_components = recommended_components + .iter() + .chain(optional_components.iter()) + .filter_map(|component| { + get_component_with_instances(store, query, ent_path, component.as_ref().into()) + .map(|(_, component_result)| component_result.into()) + }); + + Ok(ArchetypeView::from_components( + required_components.chain(other_components), + )) +} + /// Helper used to create an example store we can use for querying in doctests pub fn __populate_example_store() -> DataStore { use re_components::{datagen::build_frame_nr, ColorRGBA, Point2D}; @@ -264,3 +348,41 @@ fn simple_query_entity() { let _used = entity_view; } } + +// Minimal test matching the doctest for `query_entity_with_primary` +#[test] +fn simple_query_archetype() { + use re_arrow_store::LatestAtQuery; + use re_log_types::Timeline; + use re_types::archetypes::Points2D; + use re_types::components::{Color, Point2D}; + + let store = __populate_example_store(); + + let ent_path = "point"; + let query = LatestAtQuery::new(Timeline::new_sequence("frame_nr"), 123.into()); + + let arch_view = query_archetype::(&store, &query, &ent_path.into()).unwrap(); + + let expected_points = [Point2D::new(1.0, 2.0), Point2D::new(3.0, 4.0)]; + let expected_colors = [None, Some(Color::from_unmultiplied_rgba(255, 0, 0, 0))]; + + let view_points: Vec<_> = arch_view + .iter_required_component::() + .unwrap() + .collect(); + + let view_colors: Vec<_> = arch_view + .iter_optional_component::() + .unwrap() + .collect(); + + assert_eq!(expected_points, view_points.as_slice()); + assert_eq!(expected_colors, view_colors.as_slice()); + + #[cfg(feature = "polars")] + { + let df = arch_view.as_df2::().unwrap(); + eprintln!("{df:?}"); + } +} diff --git a/crates/re_query/src/range.rs b/crates/re_query/src/range.rs index 890bb2163dac..19399775ebc2 100644 --- a/crates/re_query/src/range.rs +++ b/crates/re_query/src/range.rs @@ -1,8 +1,9 @@ use itertools::Itertools as _; use re_arrow_store::{DataStore, LatestAtQuery, RangeQuery, TimeInt}; use re_log_types::{Component, ComponentName, EntityPath}; +use re_types::Archetype; -use crate::{get_component_with_instances, ComponentWithInstances, EntityView}; +use crate::{get_component_with_instances, ArchetypeView, ComponentWithInstances, EntityView}; // --- @@ -142,3 +143,145 @@ pub fn range_entity_with_primary<'a, Primary: Component + 'a, const N: usize>( }) }) } + +/// Iterates over the rows of any number of components and their respective cluster keys, all from +/// the point-of-view of the required components, returning an iterator of [`ArchetypeView`]s. +/// +/// An initial entity-view is yielded with the latest-at state at the start of the time range, if +/// there is any. +/// +/// The iterator only ever yields entity-views iff a required component has changed: a change +/// affecting only optional components will not yield an entity-view. +/// However, the changes in those secondary components will be accumulated into the next yielded +/// entity-view. +/// +/// This is a streaming-join: every yielded [`ArchetypeView`] will be the result of joining the latest +/// known state of all components, from their respective point-of-views. +/// +/// ⚠ The semantics are subtle! See `examples/range.rs` for an example of use. +pub fn range_archetype<'a, A: Archetype + 'a, const N: usize>( + store: &'a DataStore, + query: &RangeQuery, + ent_path: &'a EntityPath, +) -> impl Iterator, ArchetypeView)> + 'a { + re_tracing::profile_function!(); + + // TODO(jleibs) this shim is super gross + let components: [ComponentName; N] = A::all_components() + .iter() + .map(|c| c.as_ref().into()) + .collect::>() + .try_into() + .unwrap(); + + let primary: ComponentName = A::recommended_components()[0].as_ref().into(); + let cluster_key = store.cluster_key(); + + // TODO(cmc): Ideally, we'd want to simply add the cluster and primary key to the `components` + // array if they are missing, yielding either `[ComponentName; N+1]` or `[ComponentName; N+2]`. + // Unfortunately this is not supported on stable at the moment, and requires + // feature(generic_const_exprs) on nightly. + // + // The alternative to these assertions (and thus putting the burden on the caller), for now, + // would be to drop the constant sizes all the way down, which would be way more painful to + // deal with. + assert!(components.contains(&cluster_key)); + assert!(components.contains(&primary)); + + let cluster_col = components + .iter() + .find_position(|component| **component == cluster_key) + .map(|(col, _)| col) + .unwrap(); // asserted on entry + let primary_col = components + .iter() + .find_position(|component| **component == primary) + .map(|(col, _)| col) + .unwrap(); // asserted on entry + + let mut state: Vec<_> = std::iter::repeat_with(|| None) + .take(components.len()) + .collect(); + + // NOTE: This will return none for `TimeInt::Min`, i.e. range queries that start infinitely far + // into the past don't have a latest-at state! + let latest_time = query.range.min.as_i64().checked_sub(1).map(Into::into); + + let mut cwis_latest = None; + if let Some(latest_time) = latest_time { + let mut cwis_latest_raw: Vec<_> = std::iter::repeat_with(|| None) + .take(components.len()) + .collect(); + + // Fetch the latest data for every single component from their respective point-of-views, + // this will allow us to build up the initial state and send an initial latest-at + // entity-view if needed. + for (i, primary) in components.iter().enumerate() { + cwis_latest_raw[i] = get_component_with_instances( + store, + &LatestAtQuery::new(query.timeline, latest_time), + ent_path, + *primary, + ); + } + + if cwis_latest_raw[primary_col].is_some() { + cwis_latest = Some(cwis_latest_raw); + } + } + + // send the latest-at state before anything else + cwis_latest + .into_iter() + .map(move |cwis| (latest_time, true, cwis)) + .chain( + store + .range(query, ent_path, components) + .map(move |(time, row_id, mut cells)| { + // 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() + .map(|cell| { + cell.map(|cell| { + ( + row_id, + ComponentWithInstances { + instance_keys: instance_keys.clone(), /* shallow */ + values: cell, + }, + ) + }) + }) + .collect::>(); + (time, is_primary, cwis) + }), + ) + .filter_map(move |(time, is_primary, cwis)| { + for (i, cwi) in cwis + .into_iter() + .enumerate() + .filter(|(_, cwi)| cwi.is_some()) + { + state[i] = cwi; + } + + // We only yield if the primary component has been updated! + is_primary.then(|| { + let (row_id, _) = state[primary_col].clone().unwrap(); // shallow + + let components: Vec<_> = state + .clone() + .into_iter() + .filter_map(|cwi| cwi.map(|(_, cwi)| cwi.into())) + .collect(); + + let mut arch_view = ArchetypeView::from_components(components); + arch_view.row_id = row_id; + + (time, arch_view) + }) + }) +} diff --git a/crates/re_query/src/util.rs b/crates/re_query/src/util.rs index 45b61b78f0a0..2c6681e9128a 100644 --- a/crates/re_query/src/util.rs +++ b/crates/re_query/src/util.rs @@ -1,8 +1,12 @@ use re_arrow_store::{DataStore, LatestAtQuery, RangeQuery, TimeInt, TimeRange, Timeline}; use re_data_store::ExtraQueryHistory; use re_log_types::{Component, ComponentName, EntityPath}; +use re_types::Archetype; -use crate::{query_entity_with_primary, range_entity_with_primary, EntityView}; +use crate::{ + query_archetype, query_entity_with_primary, range::range_archetype, range_entity_with_primary, + ArchetypeView, EntityView, +}; /// Either dispatch to `query_entity_with_primary` or `range_entity_with_primary` /// depending on whether `ExtraQueryHistory` is set. @@ -35,3 +39,30 @@ pub fn query_primary_with_history<'a, Primary: Component + 'a, const N: usize>( Ok(itertools::Either::Right(range.map(|(_, entity)| entity))) } } + +pub fn query_archetype_with_history<'a, A: Archetype + 'a, const N: usize>( + store: &'a DataStore, + timeline: &'a Timeline, + time: &'a TimeInt, + history: &ExtraQueryHistory, + ent_path: &'a EntityPath, +) -> crate::Result> + 'a> { + let visible_history = match timeline.typ() { + re_log_types::TimeType::Time => history.nanos, + re_log_types::TimeType::Sequence => history.sequences, + }; + + if visible_history == 0 { + let latest_query = LatestAtQuery::new(*timeline, *time); + let latest = query_archetype::(store, &latest_query, ent_path)?; + + Ok(itertools::Either::Left(std::iter::once(latest))) + } else { + let min_time = *time - TimeInt::from(visible_history); + let range_query = RangeQuery::new(*timeline, TimeRange::new(min_time, *time)); + + let range = range_archetype::(store, &range_query, ent_path); + + Ok(itertools::Either::Right(range.map(|(_, entity)| entity))) + } +} diff --git a/crates/re_query/tests/archetype_visit_tests.rs b/crates/re_query/tests/archetype_visit_tests.rs new file mode 100644 index 000000000000..94188d068d77 --- /dev/null +++ b/crates/re_query/tests/archetype_visit_tests.rs @@ -0,0 +1,283 @@ +use itertools::Itertools; +use re_query::{ArchComponentWithInstances, ArchetypeView}; +use re_types::archetypes::Points2D; +use re_types::components::{Color, InstanceKey, Point2D}; + +#[test] +fn basic_single_iter() { + let instance_keys = InstanceKey::from_iter(0..2); + let points = [ + Point2D::new(1.0, 2.0), // + Point2D::new(3.0, 4.0), + ]; + + let component = ArchComponentWithInstances::from_native(instance_keys, points); + + let results = itertools::izip!( + points.into_iter(), + component.iter_values::().unwrap() + ) + .collect_vec(); + assert_eq!(results.len(), 2); + results + .iter() + .for_each(|(a, b)| assert_eq!(a, b.as_ref().unwrap())); +} + +#[test] +fn directly_joined_iter() { + let instance_keys = InstanceKey::from_iter(0..3); + + let points = [ + Point2D::new(1.0, 2.0), // + Point2D::new(3.0, 4.0), + Point2D::new(5.0, 6.0), + ]; + + let colors = [ + Color::from_rgb(0, 0, 0), // + Color::from_rgb(1, 0, 0), + Color::from_rgb(2, 0, 0), + ]; + + let points_comp = ArchComponentWithInstances::from_native(instance_keys.clone(), points); + let colors_comp = ArchComponentWithInstances::from_native(instance_keys, colors); + + let arch_view = ArchetypeView::::from_components([points_comp, colors_comp]); + + let expected_colors = [ + Some(Color::from_rgb(0, 0, 0)), + Some(Color::from_rgb(1, 0, 0)), + Some(Color::from_rgb(2, 0, 0)), + ]; + + let results = itertools::izip!( + expected_colors.iter(), + arch_view.iter_optional_component::().unwrap() + ) + .collect_vec(); + + assert_eq!(expected_colors.len(), results.len()); + results.iter().for_each(|(a, b)| assert_eq!(*a, b)); +} + +#[test] +fn joined_iter_dense_primary() { + let point_ids = InstanceKey::from_iter(0..3); + + let points = [ + Point2D::new(1.0, 2.0), // + Point2D::new(3.0, 4.0), + Point2D::new(5.0, 6.0), + ]; + + let color_ids = [ + InstanceKey(1), // + InstanceKey(2), + ]; + + let colors = [ + Color::from_rgb(1, 0, 0), // + Color::from_rgb(2, 0, 0), + ]; + + let points_comp = ArchComponentWithInstances::from_native(point_ids, points); + let colors_comp = ArchComponentWithInstances::from_native(color_ids, colors); + + let arch_view = ArchetypeView::::from_components([points_comp, colors_comp]); + + let expected_colors = vec![ + None, // + Some(Color::from_rgb(1, 0, 0)), + Some(Color::from_rgb(2, 0, 0)), + ]; + + let results = itertools::izip!( + expected_colors.iter(), + arch_view.iter_optional_component::().unwrap() + ) + .collect_vec(); + + assert_eq!(expected_colors.len(), results.len()); + results.iter().for_each(|(a, b)| assert_eq!(*a, b)); +} + +#[test] +fn joined_iter_dense_secondary() { + let point_ids = [ + InstanceKey(0), // + InstanceKey(2), + InstanceKey(4), + ]; + + let points = [ + Point2D::new(1.0, 2.0), // + Point2D::new(3.0, 4.0), + Point2D::new(5.0, 6.0), + ]; + + let color_ids = InstanceKey::from_iter(0..5); + + let colors = [ + Color::from_rgb(0, 0, 0), // + Color::from_rgb(1, 0, 0), + Color::from_rgb(2, 0, 0), + Color::from_rgb(3, 0, 0), + Color::from_rgb(4, 0, 0), + ]; + + let points_comp = ArchComponentWithInstances::from_native(point_ids, points); + let colors_comp = ArchComponentWithInstances::from_native(color_ids, colors); + + let arch_view = ArchetypeView::::from_components([points_comp, colors_comp]); + + let expected_colors = vec![ + Some(Color::from_rgb(0, 0, 0)), // + Some(Color::from_rgb(2, 0, 0)), + Some(Color::from_rgb(4, 0, 0)), + ]; + + let results = itertools::izip!( + expected_colors.iter(), + arch_view.iter_optional_component::().unwrap() + ) + .collect_vec(); + + assert_eq!(expected_colors.len(), results.len()); + results.iter().for_each(|(a, b)| assert_eq!(*a, b)); +} + +#[test] +fn complex_joined_iter() { + let point_ids = vec![ + InstanceKey(0), // + InstanceKey(17), + InstanceKey(42), + InstanceKey(96), + ]; + + let points = vec![ + Point2D::new(1.0, 2.0), // + Point2D::new(3.0, 4.0), + Point2D::new(5.0, 6.0), + Point2D::new(7.0, 8.0), + ]; + + let color_ids = vec![ + InstanceKey(17), // + InstanceKey(19), + InstanceKey(44), + InstanceKey(96), + InstanceKey(254), + ]; + + let colors = vec![ + Color::from_rgb(17, 0, 0), // + Color::from_rgb(19, 0, 0), + Color::from_rgb(44, 0, 0), + Color::from_rgb(96, 0, 0), + Color::from_rgb(254, 0, 0), + ]; + + let points_comp = ArchComponentWithInstances::from_native(point_ids, points); + let colors_comp = ArchComponentWithInstances::from_native(color_ids, colors); + + let arch_view = ArchetypeView::::from_components([points_comp, colors_comp]); + + let expected_colors = vec![ + None, + Some(Color::from_rgb(17, 0, 0)), // + None, + Some(Color::from_rgb(96, 0, 0)), + ]; + + let results = itertools::izip!( + expected_colors.iter(), + arch_view.iter_optional_component::().unwrap() + ) + .collect_vec(); + + assert_eq!(expected_colors.len(), results.len()); + results.iter().for_each(|(a, b)| assert_eq!(*a, b)); +} + +#[test] +fn single_visit() { + let instance_keys = InstanceKey::from_iter(0..4); + let points = [ + Point2D::new(1.0, 2.0), + Point2D::new(3.0, 4.0), + Point2D::new(5.0, 6.0), + Point2D::new(7.0, 8.0), + ]; + + let points_comp = ArchComponentWithInstances::from_native(instance_keys.clone(), points); + + let arch_view = ArchetypeView::::from_components([points_comp]); + + let mut instance_key_out = Vec::::new(); + let mut points_out = Vec::::new(); + + itertools::izip!( + arch_view.iter_instance_keys(), + arch_view.iter_required_component::().unwrap() + ) + .for_each(|(inst, point)| { + instance_key_out.push(inst); + points_out.push(point); + }); + + assert_eq!(instance_key_out, instance_keys); + assert_eq!(points.as_slice(), points_out.as_slice()); +} + +#[test] +fn joint_visit() { + let points = vec![ + Point2D::new(1.0, 2.0), // + Point2D::new(3.0, 4.0), + Point2D::new(5.0, 6.0), + Point2D::new(7.0, 8.0), + Point2D::new(9.0, 10.0), + ]; + + let point_ids = InstanceKey::from_iter(0..5); + + let colors = vec![ + Color::from_rgb(255, 0, 0), // + Color::from_rgb(0, 255, 0), + ]; + + let color_ids = vec![ + InstanceKey(2), // + InstanceKey(4), + ]; + + let points_comp = ArchComponentWithInstances::from_native(point_ids, points.clone()); + let colors_comp = ArchComponentWithInstances::from_native(color_ids, colors); + + let arch_view = ArchetypeView::::from_components([points_comp, colors_comp]); + + let mut points_out = Vec::::new(); + let mut colors_out = Vec::>::new(); + + itertools::izip!( + arch_view.iter_required_component::().unwrap(), + arch_view.iter_optional_component::().unwrap() + ) + .for_each(|(point, color)| { + points_out.push(point); + colors_out.push(color); + }); + + let expected_colors = vec![ + None, + None, + Some(Color::from_rgb(255, 0, 0)), + None, + Some(Color::from_rgb(0, 255, 0)), + ]; + + assert_eq!(points, points_out); + assert_eq!(expected_colors, colors_out); +} diff --git a/crates/re_space_view_spatial/Cargo.toml b/crates/re_space_view_spatial/Cargo.toml index f25dbce7b654..0d90c1f15f82 100644 --- a/crates/re_space_view_spatial/Cargo.toml +++ b/crates/re_space_view_spatial/Cargo.toml @@ -26,6 +26,7 @@ re_log_types.workspace = true re_log.workspace = true re_query.workspace = true re_renderer = { workspace = true, features = ["import-gltf"] } +re_types = { workspace = true, features = ["ecolor", "glam"] } re_tracing.workspace = true re_ui.workspace = true re_viewer_context.workspace = true diff --git a/crates/re_space_view_spatial/src/parts/entity_iterator.rs b/crates/re_space_view_spatial/src/parts/entity_iterator.rs index e48ede22e42a..9f009733acb4 100644 --- a/crates/re_space_view_spatial/src/parts/entity_iterator.rs +++ b/crates/re_space_view_spatial/src/parts/entity_iterator.rs @@ -1,6 +1,9 @@ use re_log_types::{Component, EntityPath}; -use re_query::{query_primary_with_history, EntityView, QueryError}; +use re_query::{ + query_archetype_with_history, query_primary_with_history, ArchetypeView, EntityView, QueryError, +}; use re_renderer::DepthOffset; +use re_types::Archetype; use re_viewer_context::{ ArchetypeDefinition, SpaceViewSystemExecutionError, ViewContextCollection, ViewQuery, ViewerContext, @@ -92,3 +95,72 @@ where Ok(()) } + +/// Iterates through all entity views for a given archetype. +/// +/// The callback passed in gets passed a long an [`SpatialSceneEntityContext`] which contains +/// various useful information about an entity in the context of the current scene. +pub fn process_archetype_views<'a, A, const N: usize, F>( + ctx: &mut ViewerContext<'_>, + query: &ViewQuery<'_>, + view_ctx: &ViewContextCollection, + default_depth_offset: DepthOffset, + mut fun: F, +) -> Result<(), SpaceViewSystemExecutionError> +where + A: Archetype + 'a, + F: FnMut( + &mut ViewerContext<'_>, + &EntityPath, + ArchetypeView, + &SpatialSceneEntityContext<'_>, + ) -> Result<(), QueryError>, +{ + let transforms = view_ctx.get::()?; + let depth_offsets = view_ctx.get::()?; + let annotations = view_ctx.get::()?; + let shared_render_builders = view_ctx.get::()?; + let counter = view_ctx.get::()?; + + for (ent_path, props) in query.iter_entities() { + let Some(world_from_obj) = transforms.reference_from_entity(ent_path) else { + continue; + }; + let entity_context = SpatialSceneEntityContext { + world_from_obj, + depth_offset: *depth_offsets + .per_entity + .get(&ent_path.hash()) + .unwrap_or(&default_depth_offset), + annotations: annotations.0.find(ent_path), + shared_render_builders, + highlight: query.highlights.entity_outline_mask(ent_path.hash()), + }; + + match query_archetype_with_history::( + &ctx.store_db.entity_db.data_store, + &query.timeline, + &query.latest_at, + &props.visible_history, + ent_path, + ) + .and_then(|entity_views| { + for ent_view in entity_views { + counter.num_primitives.fetch_add( + ent_view.num_instances(), + std::sync::atomic::Ordering::Relaxed, + ); + + fun(ctx, ent_path, ent_view, &entity_context)?; + } + Ok(()) + }) { + Ok(_) | Err(QueryError::PrimaryNotFound) => {} + Err(err) => { + re_log::error_once!("Unexpected error querying {ent_path:?}: {err}"); + } + } + } + + Ok(()) +} diff --git a/crates/re_space_view_spatial/src/parts/mod.rs b/crates/re_space_view_spatial/src/parts/mod.rs index f66810e84605..d981f39be01a 100644 --- a/crates/re_space_view_spatial/src/parts/mod.rs +++ b/crates/re_space_view_spatial/src/parts/mod.rs @@ -17,6 +17,8 @@ mod transform3d_arrows; pub use cameras::CamerasPart; pub use images::Image; pub use images::ImagesPart; +use re_types::components::Color; +use re_types::Archetype; pub use spatial_view_part::SpatialViewPartData; pub use transform3d_arrows::add_axis_arrows; @@ -120,6 +122,24 @@ where })) } +/// Process [`Color`] components using annotations and default colors. +pub fn process_colors_arch<'a, A: Archetype>( + arch_view: &'a re_query::ArchetypeView, + ent_path: &'a EntityPath, + annotation_infos: &'a [ResolvedAnnotationInfo], +) -> Result + 'a, re_query::QueryError> { + re_tracing::profile_function!(); + let default_color = DefaultColor::EntityPath(ent_path); + + Ok(itertools::izip!( + annotation_infos.iter(), + arch_view.iter_optional_component::()?, + ) + .map(move |(annotation_info, color)| { + annotation_info.color(color.map(move |c| c.to_array()).as_ref(), default_color) + })) +} + /// Process [`Radius`] components to [`re_renderer::Size`] using auto size where no radius is specified. pub fn process_radii<'a, Primary>( ent_path: &EntityPath, @@ -149,6 +169,33 @@ where })) } +/// Process [`Radius`] components to [`re_renderer::Size`] using auto size where no radius is specified. +pub fn process_radii_arch<'a, A: Archetype>( + arch_view: &'a re_query::ArchetypeView, + ent_path: &EntityPath, +) -> Result + 'a, re_query::QueryError> { + re_tracing::profile_function!(); + let ent_path = ent_path.clone(); + Ok(arch_view + .iter_optional_component::()? + .map(move |radius| { + radius.map_or(re_renderer::Size::AUTO, |r| { + if 0.0 <= r.0 && r.0.is_finite() { + re_renderer::Size::new_scene(r.0) + } else { + if r.0 < 0.0 { + re_log::warn_once!("Found negative radius in entity {ent_path}"); + } else if r.0.is_infinite() { + re_log::warn_once!("Found infinite radius in entity {ent_path}"); + } else { + re_log::warn_once!("Found NaN radius in entity {ent_path}"); + } + re_renderer::Size::AUTO + } + }) + })) +} + /// Resolves all annotations and keypoints for the given entity view. fn process_annotations_and_keypoints( query: &ViewQuery<'_>, @@ -198,6 +245,54 @@ where Ok((annotation_info, keypoints)) } +/// Resolves all annotations and keypoints for the given entity view. +fn process_annotations_and_keypoints_arch( + query: &ViewQuery<'_>, + arch_view: &re_query::ArchetypeView, + annotations: &Arc, +) -> Result<(Vec, Keypoints), re_query::QueryError> +where + Primary: re_types::Component + Clone, + glam::Vec3: std::convert::From, +{ + re_tracing::profile_function!(); + + let mut keypoints: Keypoints = HashMap::default(); + + // No need to process annotations if we don't have keypoints or class-ids + if !arch_view.has_component::() + && !arch_view.has_component::() + { + let resolved_annotation = annotations.class_description(None).annotation_info(); + return Ok(( + vec![resolved_annotation; arch_view.num_instances()], + keypoints, + )); + } + + let annotation_info = itertools::izip!( + arch_view.iter_required_component::()?, + arch_view.iter_optional_component::()?, + arch_view.iter_optional_component::()?, + ) + .map(|(position, keypoint_id, class_id)| { + let class_description = annotations.class_description(class_id.map(|c| c.into())); + + if let (Some(keypoint_id), Some(class_id), position) = (keypoint_id, class_id, position) { + keypoints + .entry((class_id.into(), query.latest_at.as_i64())) + .or_insert_with(Default::default) + .insert(keypoint_id.into(), position.into()); + class_description.annotation_info_with_keypoint(keypoint_id.into()) + } else { + class_description.annotation_info() + } + }) + .collect(); + + Ok((annotation_info, keypoints)) +} + #[derive(Clone)] pub enum UiLabelTarget { /// Labels a given rect (in scene coordinates) diff --git a/crates/re_space_view_spatial/src/parts/points2d.rs b/crates/re_space_view_spatial/src/parts/points2d.rs index e0404fd85e24..4a336e6d8378 100644 --- a/crates/re_space_view_spatial/src/parts/points2d.rs +++ b/crates/re_space_view_spatial/src/parts/points2d.rs @@ -1,8 +1,10 @@ -use re_components::{ - ClassId, ColorRGBA, Component, InstanceKey, KeypointId, Label, Point2D, Radius, -}; use re_data_store::{EntityPath, InstancePathHash}; -use re_query::{EntityView, QueryError}; +use re_query::{ArchetypeView, QueryError}; +use re_types::{ + archetypes::Points2D, + components::{Label, Point2D}, + Archetype, +}; use re_viewer_context::{ ArchetypeDefinition, ResolvedAnnotationInfo, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartSystem, ViewQuery, ViewerContext, @@ -11,15 +13,14 @@ use re_viewer_context::{ use crate::{ contexts::{EntityDepthOffsets, SpatialSceneEntityContext}, parts::{ - entity_iterator::process_entity_views, load_keypoint_connections, UiLabel, UiLabelTarget, + entity_iterator::process_archetype_views, load_keypoint_connections, + process_annotations_and_keypoints_arch, process_colors_arch, process_radii_arch, UiLabel, + UiLabelTarget, }, view_kind::SpatialSpaceViewKind, }; -use super::{ - picking_id_from_instance_key, process_annotations_and_keypoints, process_colors, process_radii, - SpatialViewPartData, -}; +use super::{picking_id_from_instance_key, SpatialViewPartData}; pub struct Points2DPart { /// If the number of points in the batch is > max_labels, don't render point labels. @@ -38,15 +39,15 @@ impl Default for Points2DPart { impl Points2DPart { fn process_labels<'a>( - entity_view: &'a EntityView, + arch_view: &'a ArchetypeView, instance_path_hashes: &'a [InstancePathHash], colors: &'a [egui::Color32], annotation_infos: &'a [ResolvedAnnotationInfo], ) -> Result + 'a, QueryError> { let labels = itertools::izip!( annotation_infos.iter(), - entity_view.iter_primary()?, - entity_view.iter_component::