Use ComponentDescriptor in HybridLatestAtResult (#9938)

### Related

* Part of #6889.
* Builds on #9929.
* Requires #9933.

### What

Title.

---------

Co-authored-by: Andreas Reich <[email protected]>
Co-authored-by: Andreas Reich <[email protected]>

Author: grtlr

crates/store/re_chunk/src/chunk.rs

@@ -103,10 +103,8 @@ impl ChunkComponents {
     }
 
     /// Whether any of the components in this chunk has the given name.
-    pub fn contains_component_name(&self, component_name: ComponentName) -> bool {
-        self.0
-            .keys()
-            .any(|desc| desc.component_name == component_name)
+    pub fn contains_component(&self, component_descr: &ComponentDescriptor) -> bool {
+        self.0.contains_key(component_descr)
     }
 
     /// Whether any of the components in this chunk is tagged with the given archetype name.

crates/store/re_chunk_store/src/query.rs

@@ -10,8 +10,7 @@ use re_chunk::{Chunk, LatestAtQuery, RangeQuery, TimelineName};
 use re_log_types::ResolvedTimeRange;
 use re_log_types::{EntityPath, TimeInt, Timeline};
 use re_types_core::{
-    ComponentDescriptor, ComponentDescriptorSet, ComponentName, ComponentNameSet,
-    UnorderedComponentDescriptorSet,
+    ComponentDescriptor, ComponentDescriptorSet, ComponentName, UnorderedComponentDescriptorSet,
 };
 
 use crate::{ChunkStore, store::ChunkIdSetPerTime};
@@ -146,56 +145,6 @@ impl ChunkStore {
         }
     }
 
-    /// Retrieve all the [`ComponentName`]s that have been written to for a given [`EntityPath`] on
-    /// the specified [`Timeline`].
-    ///
-    /// Static components are always included in the results.
-    ///
-    /// Returns `None` if the entity doesn't exist at all on this `timeline`.
-    // TODO(#6889): Remove in favor of `all_components_on_timeline_sorted`.
-    pub fn all_components_on_timeline_sorted_by_name(
-        &self,
-        timeline: &TimelineName,
-        entity_path: &EntityPath,
-    ) -> Option<ComponentNameSet> {
-        re_tracing::profile_function!();
-
-        let static_components: Option<ComponentNameSet> = self
-            .static_chunk_ids_per_entity
-            .get(entity_path)
-            .map(|static_chunks_per_component| {
-                static_chunks_per_component
-                    .keys()
-                    .map(|descr| descr.component_name)
-                    .collect::<ComponentNameSet>()
-            })
-            .filter(|names| !names.is_empty());
-
-        let temporal_components: Option<ComponentNameSet> = self
-            .temporal_chunk_ids_per_entity_per_component
-            .get(entity_path)
-            .map(|temporal_chunk_ids_per_timeline| {
-                temporal_chunk_ids_per_timeline
-                    .get(timeline)
-                    .map(|temporal_chunk_ids_per_component| {
-                        temporal_chunk_ids_per_component
-                            .keys()
-                            .map(|descr| descr.component_name)
-                            .collect::<ComponentNameSet>()
-                    })
-                    .unwrap_or_default()
-            })
-            .filter(|names| !names.is_empty());
-
-        match (static_components, temporal_components) {
-            (None, None) => None,
-            (None, Some(comps)) | (Some(comps), None) => Some(comps),
-            (Some(static_comps), Some(temporal_comps)) => {
-                Some(static_comps.into_iter().chain(temporal_comps).collect())
-            }
-        }
-    }
-
     /// Retrieve all the [`ComponentName`]s that have been written to for a given [`EntityPath`] on
     /// the specified [`Timeline`].
     ///
@@ -292,29 +241,24 @@ impl ChunkStore {
     pub fn all_components_for_entity_sorted(
         &self,
         entity_path: &EntityPath,
-    ) -> Option<ComponentNameSet> {
+    ) -> Option<ComponentDescriptorSet> {
         re_tracing::profile_function!();
 
-        let static_components: Option<ComponentNameSet> = self
+        let static_components: Option<ComponentDescriptorSet> = self
             .static_chunk_ids_per_entity
             .get(entity_path)
             .map(|static_chunks_per_component| {
-                static_chunks_per_component
-                    .keys()
-                    .map(|descr| descr.component_name)
-                    .collect()
+                static_chunks_per_component.keys().cloned().collect()
             });
 
-        let temporal_components: Option<ComponentNameSet> = self
+        let temporal_components: Option<ComponentDescriptorSet> = self
             .temporal_chunk_ids_per_entity_per_component
             .get(entity_path)
             .map(|temporal_chunk_ids_per_timeline| {
                 temporal_chunk_ids_per_timeline
                     .iter()
                     .flat_map(|(_, temporal_chunk_ids_per_component)| {
-                        temporal_chunk_ids_per_component
-                            .keys()
-                            .map(|descr| descr.component_name)
+                        temporal_chunk_ids_per_component.keys().cloned()
                     })
                     .collect()
             });

crates/store/re_chunk_store/tests/reads.rs

@@ -12,7 +12,7 @@ use re_log_types::{
     example_components::{MyColor, MyIndex, MyPoint},
 };
 use re_types::{
-    ComponentDescriptor, ComponentNameSet,
+    ComponentDescriptor, ComponentDescriptorSet,
     testing::{LargeStruct, build_some_large_structs},
 };
 use re_types_core::Component as _;
@@ -57,12 +57,10 @@ fn all_components() -> anyhow::Result<()> {
         |store: &ChunkStore, entity_path: &EntityPath, expected: Option<&[ComponentDescriptor]>| {
             let timeline = TimelineName::new("frame_nr");
 
-            let component_names =
-                store.all_components_on_timeline_sorted_by_name(&timeline, entity_path);
+            let component_names = store.all_components_on_timeline_sorted(&timeline, entity_path);
 
             let expected_component_names = expected.map(|expected| {
-                let expected: ComponentNameSet =
-                    expected.iter().map(|desc| desc.component_name).collect();
+                let expected: ComponentDescriptorSet = expected.iter().cloned().collect();
                 expected
             });
 

crates/store/re_query/src/latest_at.rs

@@ -436,9 +436,8 @@ impl LatestAtResults {
         &self,
         log_level: re_log::Level,
         instance_index: usize,
+        component_descr: &ComponentDescriptor,
     ) -> Option<C> {
-        let component_descr = self.find_component_descriptor(C::name())?;
-
         self.components.get(component_descr).and_then(|unit| {
             self.ok_or_log_err(
                 log_level,
@@ -452,8 +451,17 @@ impl LatestAtResults {
     ///
     /// Logs an error if the data cannot be deserialized, or if the instance index is out of bounds.
     #[inline]
-    pub fn component_instance<C: Component>(&self, instance_index: usize) -> Option<C> {
-        self.component_instance_with_log_level(re_log::Level::Error, instance_index)
+    pub fn component_instance<C: Component>(
+        &self,
+        instance_index: usize,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<C> {
+        debug_assert_eq!(component_descr.component_name, C::name());
+        self.component_instance_with_log_level(
+            re_log::Level::Error,
+            instance_index,
+            component_descr,
+        )
     }
 
     /// Returns the deserialized data for the specified component at the given instance index.

crates/store/re_types_core/src/lib.rs

@@ -44,7 +44,7 @@ pub use self::{
     component_descriptor::ComponentDescriptor,
     id::{ChunkId, RowId},
     loggable::{
-        Component, ComponentDescriptorSet, ComponentName, ComponentNameSet, DatatypeName, Loggable,
+        Component, ComponentDescriptorSet, ComponentName, DatatypeName, Loggable,
         UnorderedComponentDescriptorSet,
     },
     loggable_batch::{

crates/store/re_types_core/src/loggable.rs

@@ -123,8 +123,6 @@ pub trait Component: Loggable {
 
 pub type UnorderedComponentDescriptorSet = IntSet<ComponentDescriptor>;
 
-pub type ComponentNameSet = std::collections::BTreeSet<ComponentName>;
-
 pub type ComponentDescriptorSet = std::collections::BTreeSet<ComponentDescriptor>;
 
 re_string_interner::declare_new_type!(

crates/viewer/re_view/src/annotation_scene_context.rs

@@ -1,4 +1,4 @@
-use re_types::{Archetype as _, ComponentNameSet, archetypes::AnnotationContext};
+use re_types::{Archetype as _, ComponentDescriptorSet, archetypes::AnnotationContext};
 use re_viewer_context::{
     AnnotationMap, IdentifiedViewSystem, ViewContextSystem, ViewSystemIdentifier,
 };
@@ -13,12 +13,12 @@ impl IdentifiedViewSystem for AnnotationSceneContext {
 }
 
 impl ViewContextSystem for AnnotationSceneContext {
-    fn compatible_component_sets(&self) -> Vec<ComponentNameSet> {
+    fn compatible_component_sets(&self) -> Vec<ComponentDescriptorSet> {
         vec![
             AnnotationContext::required_components()
                 .iter()
-                .map(|descr| descr.component_name)
-                .collect(), //
+                .cloned()
+                .collect(),
         ]
     }
 

crates/viewer/re_view/src/results_ext.rs

@@ -2,7 +2,7 @@ use std::{borrow::Cow, sync::Arc};
 
 use itertools::Itertools as _;
 
-use re_chunk_store::{Chunk, LatestAtQuery, RangeQuery, UnitChunkShared};
+use re_chunk_store::{Chunk, LatestAtQuery, RangeQuery};
 use re_log_types::hash::Hash64;
 use re_query::{LatestAtResults, RangeResults};
 use re_types::ComponentDescriptor;
@@ -39,16 +39,7 @@ pub struct HybridRangeResults<'a> {
 }
 
 impl HybridLatestAtResults<'_> {
-    /// Returns the [`UnitChunkShared`] for the specified [`re_types_core::Component`].
-    #[inline]
-    pub fn get(&self, component_name: impl Into<ComponentName>) -> Option<&UnitChunkShared> {
-        let component_name = component_name.into();
-        self.overrides
-            .get_by_name(&component_name)
-            .or_else(|| self.results.get_by_name(&component_name))
-            .or_else(|| self.defaults.get_by_name(&component_name))
-    }
-
+    // TODO(#6889): Right now, fallbacks are on a per-component basis, so it's fine to pass the component name here.
     pub fn fallback_raw(&self, component_name: ComponentName) -> arrow::array::ArrayRef {
         let query_context = QueryContext {
             viewer_ctx: self.ctx.viewer_ctx,
@@ -68,41 +59,66 @@ impl HybridLatestAtResults<'_> {
 
     /// Utility for retrieving the first instance of a component, ignoring defaults.
     #[inline]
-    pub fn get_required_mono<C: re_types_core::Component>(&self) -> Option<C> {
-        self.get_required_instance(0)
+    pub fn get_required_mono<C: re_types_core::Component>(
+        &self,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<C> {
+        self.get_required_instance(0, component_descr)
     }
 
     /// Utility for retrieving the first instance of a component.
     #[inline]
-    pub fn get_mono<C: re_types_core::Component>(&self) -> Option<C> {
-        self.get_instance(0)
+    pub fn get_mono<C: re_types_core::Component>(
+        &self,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<C> {
+        self.get_instance(0, component_descr)
     }
 
     /// Utility for retrieving the first instance of a component.
     #[inline]
-    pub fn get_mono_with_fallback<C: re_types_core::Component + Default>(&self) -> C {
-        self.get_instance_with_fallback(0)
+    pub fn get_mono_with_fallback<C: re_types_core::Component + Default>(
+        &self,
+        component_descr: &ComponentDescriptor,
+    ) -> C {
+        debug_assert_eq!(component_descr.component_name, C::name());
+
+        self.get_instance_with_fallback(0, component_descr)
     }
 
     /// Utility for retrieving a single instance of a component, not checking for defaults.
     ///
     /// If overrides or defaults are present, they will only be used respectively if they have a component at the specified index.
     #[inline]
-    pub fn get_required_instance<C: re_types_core::Component>(&self, index: usize) -> Option<C> {
-        self.overrides.component_instance::<C>(index).or_else(||
+    pub fn get_required_instance<C: re_types_core::Component>(
+        &self,
+        index: usize,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<C> {
+        self.overrides
+            .component_instance::<C>(index, component_descr)
+            .or_else(||
                 // No override -> try recording store instead
-                self.results.component_instance::<C>(index))
+                self.results.component_instance::<C>(index, component_descr))
     }
 
     /// Utility for retrieving a single instance of a component.
     ///
     /// If overrides or defaults are present, they will only be used respectively if they have a component at the specified index.
     #[inline]
-    pub fn get_instance<C: re_types_core::Component>(&self, index: usize) -> Option<C> {
-        self.get_required_instance(index).or_else(|| {
-            // No override & no store -> try default instead
-            self.defaults.component_instance::<C>(index)
-        })
+    pub fn get_instance<C: re_types_core::Component>(
+        &self,
+        index: usize,
+        component_descr: &ComponentDescriptor,
+    ) -> Option<C> {
+        debug_assert_eq!(component_descr.component_name, C::name());
+
+        self.get_required_instance(index, component_descr)
+            .or_else(|| {
+                // No override & no store -> try default instead
+                self.defaults
+                    .component_instance::<C>(index, component_descr)
+            })
     }
 
     /// Utility for retrieving a single instance of a component.
@@ -112,11 +128,15 @@ impl HybridLatestAtResults<'_> {
     pub fn get_instance_with_fallback<C: re_types_core::Component + Default>(
         &self,
         index: usize,
+        component_descr: &ComponentDescriptor,
     ) -> C {
-        self.get_instance(index)
+        debug_assert_eq!(component_descr.component_name, C::name());
+
+        let component_name = component_descr.component_name;
+        self.get_instance(index, component_descr)
             .or_else(|| {
                 // No override, no store, no default -> try fallback instead
-                let raw_fallback = self.fallback_raw(C::name());
+                let raw_fallback = self.fallback_raw(component_name);
                 C::from_arrow(raw_fallback.as_ref())
                     .ok()
                     .and_then(|r| r.first().cloned())

crates/viewer/re_view_bar_chart/src/visualizer_system.rs

@@ -59,12 +59,14 @@ impl VisualizerSystem for BarChartVisualizerSystem {
             let results = data_result
                 .latest_at_with_blueprint_resolved_data::<BarChart>(ctx, &timeline_query);
 
-            let Some(tensor) = results.get_required_mono::<components::TensorData>() else {
+            let Some(tensor) =
+                results.get_required_mono::<components::TensorData>(&BarChart::descriptor_values())
+            else {
                 continue;
             };
 
             if tensor.is_vector() {
-                let color = results.get_mono_with_fallback();
+                let color = results.get_mono_with_fallback(&BarChart::descriptor_color());
                 self.charts
                     .insert(data_result.entity_path.clone(), (tensor.0.clone(), color));
             }