Make store aware of full component descriptors (#9842)

Author: Wumpf

Cargo.lock

@@ -6923,7 +6923,6 @@ dependencies = [
  "re_query",
  "re_sorbet",
  "re_tracing",
- "re_types",
  "re_types_core",
  "similar-asserts",
  "tokio",

crates/store/re_chunk/examples/latest_at.rs

@@ -12,13 +12,13 @@ fn main() -> anyhow::Result<()> {
     let query = LatestAtQuery::new(TimelineName::new("frame"), 4);
 
     // Find all relevant data for a query:
-    let chunk = chunk.latest_at(&query, MyPoint::name());
-    eprintln!("{:?} @ {query:?}:\n{chunk}", MyPoint::name());
+    let chunk = chunk.latest_at(&query, &MyPoint::descriptor());
+    eprintln!("{:?} @ {query:?}:\n{chunk}", MyPoint::descriptor());
 
     // And then slice it as appropriate:
     let chunk = chunk
         .timeline_sliced(TimelineName::log_time())
-        .component_sliced(MyPoint::name());
+        .component_sliced(&MyPoint::descriptor());
     eprintln!("Sliced down to specific timeline and component:\n{chunk}");
 
     Ok(())

crates/store/re_chunk/examples/range.rs

@@ -15,13 +15,13 @@ fn main() -> anyhow::Result<()> {
     let query = RangeQuery::new(TimelineName::new("frame"), ResolvedTimeRange::EVERYTHING);
 
     // Find all relevant data for a query:
-    let chunk = chunk.range(&query, MyPoint::name());
-    eprintln!("{:?} @ {query:?}:\n{chunk}", MyPoint::name());
+    let chunk = chunk.range(&query, &MyPoint::descriptor());
+    eprintln!("{:?} @ {query:?}:\n{chunk}", MyPoint::descriptor());
 
     // And then slice it as appropriate:
     let chunk = chunk
         .timeline_sliced(TimelineName::log_time())
-        .component_sliced(MyPoint::name());
+        .component_sliced(&MyPoint::descriptor());
     eprintln!("Sliced down to specific timeline and component:\n{chunk}");
 
     Ok(())

crates/store/re_chunk/src/chunk.rs

@@ -221,14 +221,30 @@ impl Chunk {
     ///
     /// This is undefined behavior if there are more than one component with that name.
     //
-    // TODO(cmc): Kinda disgusting but it makes our lives easier during the interim, as long as we're
+    // TODO(#6889): Kinda disgusting but it makes our lives easier during the interim, as long as we're
     // in this weird halfway in-between state where we still have a bunch of things indexed by name only.
     #[inline]
     pub fn get_first_component(&self, component_name: ComponentName) -> Option<&ArrowListArray> {
         self.components.iter().find_map(move |(descr, array)| {
             (descr.component_name == component_name).then_some(array)
         })
     }
+
+    /// Returns any component descriptor with the given [`ComponentName`].
+    ///
+    /// This is undefined behavior if there are more than one component with that name.
+    //
+    // TODO(#6889): Kinda disgusting but it makes our lives easier during the interim, as long as we're
+    // in this weird halfway in-between state where we still have a bunch of things indexed by name only.
+    #[inline]
+    pub fn get_first_component_descriptor(
+        &self,
+        component_name: ComponentName,
+    ) -> Option<&ComponentDescriptor> {
+        self.components
+            .keys()
+            .find(|descr| descr.component_name == component_name)
+    }
 }
 
 impl Chunk {
@@ -582,9 +598,12 @@ impl Chunk {
     //
     // TODO(cmc): This needs to be stored in chunk metadata and transported across IPC.
     #[inline]
-    pub fn num_events_for_component(&self, component_name: ComponentName) -> Option<u64> {
+    pub fn num_events_for_component(
+        &self,
+        component_descriptor: &ComponentDescriptor,
+    ) -> Option<u64> {
         // Reminder: component columns are sparse, we must check validity bitmap.
-        self.get_first_component(component_name).map(|list_array| {
+        self.components.get(component_descriptor).map(|list_array| {
             list_array.nulls().map_or_else(
                 || list_array.len() as u64,
                 |validity| validity.len() as u64 - validity.null_count() as u64,

crates/store/re_chunk/src/helpers.rs

@@ -4,6 +4,7 @@ use arrow::array::Array as _;
 use arrow::array::ArrayRef as ArrowArrayRef;
 
 use re_log_types::{TimeInt, TimelineName};
+use re_types_core::ComponentDescriptor;
 use re_types_core::{Component, ComponentName};
 
 use crate::{Chunk, ChunkResult, RowId};
@@ -258,20 +259,45 @@ impl UnitChunkShared {
 impl UnitChunkShared {
     // --- Batch ---
 
+    /// Returns the raw data for the specified component name.
+    ///
+    /// Indetermined which batch is returned if there are multiple components with the same name.
+    /// TODO(#6889): Can we remove this method?
+    #[inline]
+    pub fn component_batch_raw_by_component_name(
+        &self,
+        component_name: ComponentName,
+    ) -> Option<ArrowArrayRef> {
+        debug_assert!(self.num_rows() == 1);
+        self.components
+            .get_by_component_name(component_name)
+            .next()
+            .and_then(|list_array| list_array.is_valid(0).then(|| list_array.value(0)))
+    }
+
     /// Returns the raw data for the specified component.
     #[inline]
-    pub fn component_batch_raw(&self, component_name: &ComponentName) -> Option<ArrowArrayRef> {
+    pub fn component_batch_raw(
+        &self,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<ArrowArrayRef> {
         debug_assert!(self.num_rows() == 1);
-        self.get_first_component(*component_name)
+        self.components
+            .get(component_descr)
             .and_then(|list_array| list_array.is_valid(0).then(|| list_array.value(0)))
     }
 
     /// Returns the deserialized data for the specified component.
     ///
     /// Returns an error if the data cannot be deserialized.
+    /// In debug builds, panics if the descriptor doesn't have the same component name as the component type.
     #[inline]
-    pub fn component_batch<C: Component>(&self) -> Option<ChunkResult<Vec<C>>> {
-        let data = C::from_arrow(&*self.component_batch_raw(&C::name())?);
+    pub fn component_batch<C: Component>(
+        &self,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<ChunkResult<Vec<C>>> {
+        debug_assert_eq!(C::name(), component_descr.component_name);
+        let data = C::from_arrow(&*self.component_batch_raw(component_descr)?);
         Some(data.map_err(Into::into))
     }
 
@@ -283,10 +309,10 @@ impl UnitChunkShared {
     #[inline]
     pub fn component_instance_raw(
         &self,
-        component_name: &ComponentName,
+        component_descr: &ComponentDescriptor,
         instance_index: usize,
     ) -> Option<ChunkResult<ArrowArrayRef>> {
-        let array = self.component_batch_raw(component_name)?;
+        let array = self.component_batch_raw(component_descr)?;
         if array.len() > instance_index {
             Some(Ok(array.slice(instance_index, 1)))
         } else {
@@ -301,12 +327,15 @@ impl UnitChunkShared {
     /// Returns the deserialized data for the specified component at the given instance index.
     ///
     /// Returns an error if the data cannot be deserialized, or if the instance index is out of bounds.
+    /// In debug builds, panics if the descriptor doesn't have the same component name as the component type.
     #[inline]
     pub fn component_instance<C: Component>(
         &self,
+        component_descr: &ComponentDescriptor,
         instance_index: usize,
     ) -> Option<ChunkResult<C>> {
-        let res = self.component_instance_raw(&C::name(), instance_index)?;
+        debug_assert_eq!(C::name(), component_descr.component_name);
+        let res = self.component_instance_raw(component_descr, instance_index)?;
 
         let array = match res {
             Ok(array) => array,
@@ -327,9 +356,9 @@ impl UnitChunkShared {
     #[inline]
     pub fn component_mono_raw(
         &self,
-        component_name: &ComponentName,
+        component_descr: &ComponentDescriptor,
     ) -> Option<ChunkResult<ArrowArrayRef>> {
-        let array = self.component_batch_raw(component_name)?;
+        let array = self.component_batch_raw(component_descr)?;
         if array.len() == 1 {
             Some(Ok(array.slice(0, 1)))
         } else {
@@ -344,9 +373,14 @@ impl UnitChunkShared {
     /// Returns the deserialized data for the specified component, assuming a mono-batch.
     ///
     /// Returns an error if the data cannot be deserialized, or if the underlying batch is not of unit length.
+    /// In debug builds, panics if the descriptor doesn't have the same component name as the component type.
     #[inline]
-    pub fn component_mono<C: Component>(&self) -> Option<ChunkResult<C>> {
-        let res = self.component_mono_raw(&C::name())?;
+    pub fn component_mono<C: Component>(
+        &self,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<ChunkResult<C>> {
+        debug_assert_eq!(C::name(), component_descr.component_name);
+        let res = self.component_mono_raw(component_descr)?;
 
         let array = match res {
             Ok(array) => array,

crates/store/re_chunk/src/latest_at.rs

@@ -1,7 +1,7 @@
 use arrow::array::Array as _;
 
 use re_log_types::{TimeInt, TimelineName};
-use re_types_core::ComponentName;
+use re_types_core::ComponentDescriptor;
 
 use crate::{Chunk, RowId};
 
@@ -71,14 +71,14 @@ impl Chunk {
     /// information by inspecting the data, for examples timestamps on other timelines.
     /// See [`Self::timeline_sliced`] and [`Self::component_sliced`] if you do want to filter this
     /// extra data.
-    pub fn latest_at(&self, query: &LatestAtQuery, component_name: ComponentName) -> Self {
+    pub fn latest_at(&self, query: &LatestAtQuery, component_descr: &ComponentDescriptor) -> Self {
         if self.is_empty() {
             return self.clone();
         }
 
         re_tracing::profile_function!(format!("{query:?}"));
 
-        let Some(component_list_array) = self.get_first_component(component_name) else {
+        let Some(component_list_array) = self.components.get(component_descr) else {
             return self.emptied();
         };
 

crates/store/re_chunk/src/range.rs

@@ -1,5 +1,5 @@
 use re_log_types::{ResolvedTimeRange, TimeInt, TimelineName};
-use re_types_core::ComponentName;
+use re_types_core::ComponentDescriptor;
 
 use crate::Chunk;
 
@@ -193,7 +193,7 @@ impl Chunk {
     //
     // TODO(apache/arrow-rs#5375): Since we don't have access to arrow's ListView yet, we must actually clone the
     // data if the chunk requires sorting.
-    pub fn range(&self, query: &RangeQuery, component_name: ComponentName) -> Self {
+    pub fn range(&self, query: &RangeQuery, component_descr: &ComponentDescriptor) -> Self {
         if self.is_empty() {
             return self.clone();
         }
@@ -215,7 +215,7 @@ impl Chunk {
             Cow::Borrowed(self)
         };
         let chunk = if !keep_extra_components {
-            Cow::Owned(chunk.component_sliced(component_name))
+            Cow::Owned(chunk.component_sliced(component_descr))
         } else {
             chunk
         };
@@ -226,7 +226,7 @@ impl Chunk {
             // equivalent to just running a latest-at query.
             chunk.latest_at(
                 &crate::LatestAtQuery::new(*query.timeline(), TimeInt::MAX),
-                component_name,
+                component_descr,
             )
         } else {
             let Some(is_sorted_by_time) = chunk
@@ -237,7 +237,7 @@ impl Chunk {
                 return chunk.emptied();
             };
 
-            let chunk = chunk.densified(component_name);
+            let chunk = chunk.densified(component_descr);
 
             let chunk = if is_sorted_by_time {
                 // Temporal, row-sorted, time-sorted chunk

crates/store/re_chunk/src/slice.rs

@@ -168,18 +168,18 @@ impl Chunk {
         chunk
     }
 
-    /// Slices the [`Chunk`] horizontally by keeping only the selected `component_name`.
+    /// Slices the [`Chunk`] horizontally by keeping only the selected `component_descr`.
     ///
     /// The result is a new [`Chunk`] with the same rows and (at-most) one component column.
     /// All non-component columns will be kept as-is.
     ///
-    /// If `component_name` is not found within the [`Chunk`], the end result will be the same as the
+    /// If `component_descr` is not found within the [`Chunk`], the end result will be the same as the
     /// current chunk but without any component column.
     ///
     /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
     #[must_use]
     #[inline]
-    pub fn component_sliced(&self, component_name: ComponentName) -> Self {
+    pub fn component_sliced(&self, component_descr: &ComponentDescriptor) -> Self {
         let Self {
             id,
             entity_path,
@@ -199,9 +199,9 @@ impl Chunk {
             timelines: timelines.clone(),
             components: crate::ChunkComponents(
                 components
-                    .iter()
-                    .filter(|&(desc, _list_array)| (desc.component_name == component_name))
-                    .map(|(desc, list_array)| (desc.clone(), list_array.clone()))
+                    .get(component_descr)
+                    .map(|list_array| (component_descr.clone(), list_array.clone()))
+                    .into_iter()
                     .collect(),
             ),
         };
@@ -305,20 +305,20 @@ impl Chunk {
         chunk
     }
 
-    /// Densifies the [`Chunk`] vertically based on the `component_name` column.
+    /// Densifies the [`Chunk`] vertically based on the `component_descriptor` column.
     ///
-    /// Densifying here means dropping all rows where the associated value in the `component_name`
+    /// Densifying here means dropping all rows where the associated value in the `component_descriptor`
     /// column is null.
     ///
-    /// The result is a new [`Chunk`] where the `component_name` column is guaranteed to be dense.
+    /// The result is a new [`Chunk`] where the `component_descriptor` column is guaranteed to be dense.
     ///
-    /// If `component_name` doesn't exist in this [`Chunk`], or if it is already dense, this method
+    /// If `component_descriptor` doesn't exist in this [`Chunk`], or if it is already dense, this method
     /// is a no-op.
     ///
     /// WARNING: the returned chunk has the same old [`crate::ChunkId`]! Change it with [`Self::with_id`].
     #[must_use]
     #[inline]
-    pub fn densified(&self, component_name_pov: ComponentName) -> Self {
+    pub fn densified(&self, component_descr_pov: &ComponentDescriptor) -> Self {
         let Self {
             id,
             entity_path,
@@ -333,7 +333,7 @@ impl Chunk {
             return self.clone();
         }
 
-        let Some(component_list_array) = self.get_first_component(component_name_pov) else {
+        let Some(component_list_array) = self.components.get(component_descr_pov) else {
             return self.clone();
         };
 
@@ -361,7 +361,7 @@ impl Chunk {
                 .iter()
                 .map(|(component_desc, list_array)| {
                     let filtered = re_arrow_util::filter_array(list_array, &validity_filter);
-                    let filtered = if component_desc.component_name == component_name_pov {
+                    let filtered = if component_desc == component_descr_pov {
                         // Make sure we fully remove the validity bitmap for the densified
                         // component.
                         // This will allow further operations on this densified chunk to take some

crates/store/re_chunk/tests/latest_at.rs

@@ -478,7 +478,7 @@ fn query_and_compare(
 ) {
     re_log::setup_logging();
 
-    let results = chunk.latest_at(query, component_desc.component_name);
+    let results = chunk.latest_at(query, &component_desc);
 
     eprintln!("Query: {component_desc} @ {query:?}");
     eprintln!("Data:\n{chunk}");

crates/store/re_chunk/tests/range.rs

@@ -447,7 +447,7 @@ fn query_and_compare(
 ) {
     re_log::setup_logging();
 
-    let results = chunk.range(query, component_desc.component_name);
+    let results = chunk.range(query, &component_desc);
 
     eprintln!("Query: {component_desc} @ {query:?}");
     eprintln!("Data:\n{chunk}");