Speed up handling of large numbers of transform entities

crates/viewer/re_space_view_spatial/src/contexts/transform_context.rs

@@ -530,27 +530,17 @@ fn query_and_resolve_tree_transform_at_entity(
     entity_db: &EntityDb,
     query: &LatestAtQuery,
 ) -> Option<glam::Affine3A> {
-    if !TransformComponentTracker::access(entity_db.store_id(), |tracker| {
-        tracker.is_potentially_transformed_transform3d(entity_path)
-    })
-    .unwrap_or(false)
-    {
+    let Some(transform3d_components) =
+        TransformComponentTracker::access(entity_db.store_id(), |tracker| {
+            tracker.transform3d_components(entity_path).cloned()
+        })
+        .flatten()
+    else {
         return None;
-    }
+    };
 
     // TODO(#6743): Doesn't take into account overrides.
-    let result = entity_db.latest_at(
-        query,
-        entity_path,
-        [
-            Translation3D::name(),
-            RotationAxisAngle::name(),
-            RotationQuat::name(),
-            Scale3D::name(),
-            TransformMat3x3::name(),
-            TransformRelation::name(),
-        ],
-    );
+    let result = entity_db.latest_at(query, entity_path, transform3d_components);
     if result.components.is_empty() {
         return None;
     }
@@ -586,26 +576,17 @@ fn query_and_resolve_instance_poses_at_entity(
     entity_db: &EntityDb,
     query: &LatestAtQuery,
 ) -> Vec<glam::Affine3A> {
-    if !TransformComponentTracker::access(entity_db.store_id(), |tracker| {
-        tracker.is_potentially_transformed_pose3d(entity_path)
-    })
-    .unwrap_or(false)
-    {
+    let Some(pose3d_components) =
+        TransformComponentTracker::access(entity_db.store_id(), |tracker| {
+            tracker.pose3d_components(entity_path).cloned()
+        })
+        .flatten()
+    else {
         return Vec::new();
-    }
+    };
 
     // TODO(#6743): Doesn't take into account overrides.
-    let result = entity_db.latest_at(
-        query,
-        entity_path,
-        [
-            PoseTranslation3D::name(),
-            PoseRotationAxisAngle::name(),
-            PoseRotationQuat::name(),
-            PoseScale3D::name(),
-            PoseTransformMat3x3::name(),
-        ],
-    );
+    let result = entity_db.latest_at(query, entity_path, pose3d_components);
 
     let max_count = result
         .components

crates/viewer/re_space_view_spatial/src/transform_component_tracker.rs

@@ -1,7 +1,7 @@
 use ahash::HashMap;
 use once_cell::sync::OnceCell;
 
-use nohash_hasher::IntSet;
+use nohash_hasher::{IntMap, IntSet};
 use re_chunk_store::{
     ChunkStore, ChunkStoreDiffKind, ChunkStoreEvent, ChunkStoreSubscriber,
     ChunkStoreSubscriberHandle,
@@ -20,11 +20,13 @@ use re_types::ComponentName;
 /// This is a huge performance improvement in practice, especially in recordings with many entities.
 #[derive(Default)]
 pub struct TransformComponentTracker {
-    /// Which entities have had any `Transform3D` component at any point in time.
-    transform3d_entities: IntSet<EntityPathHash>,
+    /// Which entities have had any `Transform3D` component at any point in time, and which
+    /// components they actually make use of.
+    transform3d_entities: IntMap<EntityPathHash, IntSet<ComponentName>>,
 
-    /// Which entities have had any `InstancePoses3D` components at any point in time.
-    pose3d_entities: IntSet<EntityPathHash>,
+    /// Which entities have had any `InstancePoses3D` components at any point in time, and
+    /// which components they actually make use of.
+    pose3d_entities: IntMap<EntityPathHash, IntSet<ComponentName>>,
 }
 
 impl TransformComponentTracker {
@@ -41,13 +43,16 @@ impl TransformComponentTracker {
     }
 
     #[inline]
-    pub fn is_potentially_transformed_transform3d(&self, entity_path: &EntityPath) -> bool {
-        self.transform3d_entities.contains(&entity_path.hash())
+    pub fn transform3d_components(
+        &self,
+        entity_path: &EntityPath,
+    ) -> Option<&IntSet<ComponentName>> {
+        self.transform3d_entities.get(&entity_path.hash())
     }
 
     #[inline]
-    pub fn is_potentially_transformed_pose3d(&self, entity_path: &EntityPath) -> bool {
-        self.pose3d_entities.contains(&entity_path.hash())
+    pub fn pose3d_components(&self, entity_path: &EntityPath) -> Option<&IntSet<ComponentName>> {
+        self.pose3d_entities.get(&entity_path.hash())
     }
 }
 
@@ -119,15 +124,35 @@ impl ChunkStoreSubscriber for TransformComponentTrackerStoreSubscriber {
             let entity_path_hash = event.chunk.entity_path().hash();
 
             for component_name in event.chunk.component_names() {
-                if self.transform_components.contains(&component_name) {
+                if self.transform_components.contains(&component_name)
+                    && event
+                        .chunk
+                        .components()
+                        .get(&component_name)
+                        .map_or(false, |list_array| {
+                            list_array.offsets().lengths().any(|len| len > 0)
+                        })
+                {
                     transform_component_tracker
                         .transform3d_entities
-                        .insert(entity_path_hash);
+                        .entry(entity_path_hash)
+                        .or_default()
+                        .insert(component_name);
                 }
-                if self.pose_components.contains(&component_name) {
+                if self.pose_components.contains(&component_name)
+                    && event
+                        .chunk
+                        .components()
+                        .get(&component_name)
+                        .map_or(false, |list_array| {
+                            list_array.offsets().lengths().any(|len| len > 0)
+                        })
+                {
                     transform_component_tracker
                         .pose3d_entities
-                        .insert(entity_path_hash);
+                        .entry(entity_path_hash)
+                        .or_default()
+                        .insert(component_name);
                 }
             }
         }

crates/viewer/re_space_view_spatial/src/visualizers/transform3d_arrows.rs

@@ -1,7 +1,7 @@
 use egui::Color32;
 use nohash_hasher::IntSet;
 use re_log_types::{EntityPath, Instance};
-use re_space_view::DataResultQuery;
+use re_space_view::{latest_at_with_blueprint_resolved_data, DataResultQuery};
 use re_types::{
     archetypes::{Pinhole, Transform3D},
     components::{AxisLength, ImagePlaneDistance},
@@ -122,8 +122,18 @@ impl VisualizerSystem for Transform3DArrowsVisualizer {
                     .single_entity_transform_required(&data_result.entity_path, "Transform3DArrows")
             };
 
-            let results = data_result
-                .latest_at_with_blueprint_resolved_data::<Transform3D>(ctx, &latest_at_query);
+            // Note, we use this interface instead of `data_result.latest_at_with_blueprint_resolved_data` to avoid querying
+            // for a bunch of unused components. The actual transform data comes out of the context manager and can't be
+            // overridden via blueprint anyways.
+            let results = latest_at_with_blueprint_resolved_data(
+                ctx,
+                None,
+                &latest_at_query,
+                data_result,
+                std::iter::once(AxisLength::name()),
+                false,
+            );
+
             let axis_length: f32 = results.get_mono_with_fallback::<AxisLength>().into();
 
             if axis_length == 0.0 {