Significantly improve transform ingestion speed (#11655)

Author: Wumpf

Cargo.lock

@@ -9651,6 +9651,7 @@ version = "0.27.0-alpha.1+dev"
 dependencies = [
  "ahash",
  "bitflags 2.9.4",
+ "criterion",
  "glam",
  "insta",
  "itertools 0.14.0",
@@ -11935,6 +11936,17 @@ dependencies = [
  "rerun",
 ]
 
+[[package]]
+name = "test_out_of_order_transforms"
+version = "0.27.0-alpha.1+dev"
+dependencies = [
+ "anyhow",
+ "clap",
+ "glam",
+ "re_log",
+ "rerun",
+]
+
 [[package]]
 name = "test_ui_wakeup"
 version = "0.27.0-alpha.1+dev"

crates/store/re_grpc_client/src/write.rs

@@ -391,7 +391,7 @@ async fn message_proxy_client(
                         Some(Cmd::Flush { on_done }) => {
                             // Messages are received in order, so once we receive a `flush`
                             // we know we've sent all messages before that flush through already.
-                            re_log::debug!("Flush requested");
+                            re_log::trace!("Flush requested");
                             if on_done.send(()).is_err() {
                                 // Flush channel may already be closed for non-blocking flush, so this isn't an error.
                                 re_log::debug!("Failed to respond to flush: flush report channel was closed");

crates/store/re_tf/Cargo.toml

@@ -36,4 +36,12 @@ thiserror.workspace = true
 vec1 = { workspace = true, features = ["smallvec-v1"] }
 
 [dev-dependencies]
+criterion.workspace = true
 insta.workspace = true
+
+[lib]
+bench = false
+
+[[bench]]
+name = "transform_resolution_cache_bench"
+harness = false

crates/store/re_tf/benches/transform_resolution_cache_bench.rs

@@ -0,0 +1,116 @@
+#![expect(clippy::unwrap_used)] // acceptable in benchmarks
+
+use criterion::{Criterion, criterion_group, criterion_main};
+use itertools::Itertools as _;
+use std::sync::Arc;
+
+use re_chunk_store::{Chunk, ChunkStoreEvent};
+use re_entity_db::EntityDb;
+use re_log_types::{EntityPath, StoreId, TimePoint, Timeline, TimelineName};
+use re_types::{RowId, archetypes};
+
+use re_tf::{TransformFrameIdHash, TransformResolutionCache};
+
+const NUM_TIMELINES: usize = 4;
+const NUM_TIMEPOINTS: usize = 1000;
+const NUM_TIMEPOINTS_PER_ENTITY: usize = 50;
+const NUM_ENTITIES: usize = 100;
+
+fn setup_store() -> (EntityDb, Vec<ChunkStoreEvent>) {
+    let mut entity_db = EntityDb::new(StoreId::random(
+        re_log_types::StoreKind::Recording,
+        "test_app",
+    ));
+
+    let timelines = (0..NUM_TIMELINES)
+        .map(|i| Timeline::new(format!("timeline{i}"), re_log_types::TimeType::Sequence))
+        .collect_vec();
+
+    let mut events = Vec::new();
+    for entity_idx in 0..NUM_ENTITIES {
+        for batch in 0..(NUM_TIMEPOINTS / NUM_TIMEPOINTS_PER_ENTITY) {
+            let chunk_base_time = batch * NUM_TIMEPOINTS_PER_ENTITY;
+
+            let mut builder = Chunk::builder(EntityPath::from(format!("entity{entity_idx}")));
+            for t in 0..NUM_TIMEPOINTS_PER_ENTITY {
+                let mut timepoint = TimePoint::default();
+                for timeline in &timelines {
+                    #[expect(clippy::cast_possible_wrap)]
+                    timepoint.insert(*timeline, (chunk_base_time + t) as i64);
+                }
+                builder = builder.with_archetype(
+                    RowId::new(),
+                    timepoint,
+                    &archetypes::Transform3D::from_translation([1.0, 2.0, 3.0])
+                        .with_scale(2.0)
+                        .with_quaternion(glam::Quat::from_xyzw(0.0, 2.0, 3.0, 1.0))
+                        .with_mat3x3(glam::Mat3::IDENTITY),
+                );
+            }
+            let chunk = builder.build().unwrap();
+
+            events.extend(entity_db.add_chunk(&Arc::new(chunk)).unwrap().into_iter());
+        }
+    }
+    (entity_db, events)
+}
+
+fn transform_resolution_cache_query(c: &mut Criterion) {
+    let (entity_db, events) = setup_store();
+
+    c.bench_function("build_from_entitydb", |b| {
+        b.iter(|| {
+            let mut cache = TransformResolutionCache::default();
+            cache.process_store_events(events.iter());
+            cache
+        });
+    });
+
+    let query = re_chunk_store::LatestAtQuery::new(TimelineName::new("timeline2"), 123);
+    let queried_frame = TransformFrameIdHash::from_entity_path(&EntityPath::from("entity2"));
+
+    c.bench_function("query_uncached_frame", |b| {
+        b.iter_batched(
+            || {
+                let mut cache = TransformResolutionCache::default();
+                cache.process_store_events(events.iter());
+                cache
+            },
+            |mut cold_cache| {
+                let frame_transforms = cold_cache
+                    .transforms_for_timeline(query.timeline())
+                    .frame_transforms(queried_frame)
+                    .unwrap();
+                frame_transforms
+                    .latest_at_transform(&entity_db, &query)
+                    .unwrap()
+            },
+            criterion::BatchSize::PerIteration,
+        );
+    });
+
+    let mut warm_cache = TransformResolutionCache::default();
+    warm_cache.process_store_events(events.iter());
+    warm_cache
+        .transforms_for_timeline(query.timeline())
+        .frame_transforms(queried_frame)
+        .unwrap()
+        .latest_at_transform(&entity_db, &query);
+
+    c.bench_function("query_cached_frame", |b| {
+        b.iter(|| {
+            let frame_transforms = warm_cache
+                .transforms_for_timeline(query.timeline())
+                .frame_transforms(queried_frame)
+                .unwrap();
+            frame_transforms
+                .latest_at_transform(&entity_db, &query)
+                .unwrap()
+        });
+    });
+
+    // TODO(andreas): Additional benchmarks for iterative invalidation would be great!
+}
+
+criterion_group!(benches, transform_resolution_cache_query);
+criterion_main!(benches);

crates/store/re_tf/src/transform_forest.rs

@@ -3,7 +3,7 @@ use nohash_hasher::IntMap;
 use vec1::smallvec_v1::SmallVec1;
 
 use re_chunk_store::LatestAtQuery;
-use re_entity_db::{EntityPath, EntityTree};
+use re_entity_db::{EntityDb, EntityPath, EntityTree};
 use re_types::ArchetypeName;
 
 use crate::{
@@ -239,7 +239,7 @@ impl TransformForest {
     /// entities are transformed relative to it.
     pub fn new(
         recording: &re_entity_db::EntityDb,
-        transform_cache: &TransformResolutionCache,
+        transform_cache: &mut TransformResolutionCache,
         time_query: &LatestAtQuery,
     ) -> Self {
         re_tracing::profile_function!();
@@ -256,6 +256,7 @@ impl TransformForest {
         };
         tree.entity_tree_gather_transforms_recursive(
             entity_tree,
+            recording,
             time_query,
             // Ignore potential pinhole camera at the root of the view, since it is regarded as being "above" this root.
             // TODO(andreas): Should we warn about that?
@@ -271,9 +272,10 @@ impl TransformForest {
     fn entity_tree_gather_transforms_recursive(
         &mut self,
         subtree: &EntityTree,
+        entity_db: &EntityDb,
         query: &LatestAtQuery,
         transform_root_from_parent: TransformInfo,
-        transforms_for_timeline: &CachedTransformsForTimeline,
+        transforms_for_timeline: &mut CachedTransformsForTimeline,
     ) {
         let root = transform_root_from_parent.root;
         let root_from_parent = transform_root_from_parent.target_from_source;
@@ -287,7 +289,8 @@ impl TransformForest {
         for child_tree in subtree.children.values() {
             let child_path = &child_tree.path;
 
-            let transforms_at_entity = transforms_at(child_path, query, transforms_for_timeline);
+            let transforms_at_entity =
+                transforms_at(child_path, entity_db, query, transforms_for_timeline);
 
             let root_from_child =
                 root_from_parent * transforms_at_entity.parent_from_entity_tree_transform;
@@ -313,11 +316,11 @@ impl TransformForest {
             // Collect & compute poses.
             let root_from_instances = compute_root_from_instances(
                 target_from_source,
-                transforms_at_entity.entity_from_instance_poses,
+                transforms_at_entity.entity_from_instance_poses.as_ref(),
             );
             let root_from_archetype = compute_root_from_archetype(
                 target_from_source,
-                transforms_at_entity.entity_from_instance_poses,
+                transforms_at_entity.entity_from_instance_poses.as_ref(),
             );
 
             let transform_root_from_child = TransformInfo {
@@ -329,6 +332,7 @@ impl TransformForest {
 
             self.entity_tree_gather_transforms_recursive(
                 child_tree,
+                entity_db,
                 query,
                 transform_root_from_child,
                 transforms_for_timeline,
@@ -656,17 +660,18 @@ fn pinhole3d_from_image_plane(
 
 /// Resolved transforms at an entity.
 #[derive(Default)]
-struct TransformsAtEntity<'a> {
+struct TransformsAtEntity {
     parent_from_entity_tree_transform: glam::Affine3A,
-    entity_from_instance_poses: Option<&'a PoseTransformArchetypeMap>,
-    pinhole_projection: Option<&'a ResolvedPinholeProjection>,
+    entity_from_instance_poses: Option<PoseTransformArchetypeMap>,
+    pinhole_projection: Option<ResolvedPinholeProjection>,
 }
 
-fn transforms_at<'a>(
+fn transforms_at(
     entity_path: &EntityPath,
+    entity_db: &EntityDb,
     query: &LatestAtQuery,
-    transforms_for_timeline: &'a CachedTransformsForTimeline,
-) -> TransformsAtEntity<'a> {
+    transforms_for_timeline: &mut CachedTransformsForTimeline,
+) -> TransformsAtEntity {
     // This is called very frequently, don't put a profile scope here.
 
     let Some(entity_transforms) = transforms_for_timeline
@@ -676,13 +681,17 @@ fn transforms_at<'a>(
     };
 
     let parent_from_entity_tree_transform = entity_transforms
-        .latest_at_transform(query)
+        .latest_at_transform(entity_db, query)
         //  TODO(RR-2511): Don't ignore target frame.
         .map_or(glam::Affine3A::IDENTITY, |source_to_target| {
             source_to_target.transform
         });
-    let entity_from_instance_poses = entity_transforms.latest_at_instance_poses_all(query);
-    let pinhole_projection = entity_transforms.latest_at_pinhole(query);
+    let entity_from_instance_poses = entity_transforms
+        .latest_at_instance_poses(entity_db, query)
+        .cloned();
+    let pinhole_projection = entity_transforms
+        .latest_at_pinhole(entity_db, query)
+        .cloned();
 
     TransformsAtEntity {
         parent_from_entity_tree_transform,
@@ -800,13 +809,10 @@ mod tests {
     fn test_simple_entity_hierarchy() {
         let test_scene = entity_hierarchy_test_scene();
         let mut transform_cache = TransformResolutionCache::default();
-        transform_cache.add_chunks(
-            &test_scene,
-            test_scene.storage_engine().store().iter_chunks(),
-        );
+        transform_cache.add_chunks(test_scene.storage_engine().store().iter_chunks());
 
         let query = LatestAtQuery::latest(TimelineName::log_tick());
-        let transform_forest = TransformForest::new(&test_scene, &transform_cache, &query);
+        let transform_forest = TransformForest::new(&test_scene, &mut transform_cache, &query);
 
         let all_entity_paths = test_scene
             .entity_paths()
@@ -935,10 +941,10 @@ mod tests {
             .unwrap();
 
         let mut transform_cache = TransformResolutionCache::default();
-        transform_cache.add_chunks(&entity_db, entity_db.storage_engine().store().iter_chunks());
+        transform_cache.add_chunks(entity_db.storage_engine().store().iter_chunks());
 
         let query = LatestAtQuery::latest(TimelineName::log_tick());
-        let transform_forest = TransformForest::new(&entity_db, &transform_cache, &query);
+        let transform_forest = TransformForest::new(&entity_db, &mut transform_cache, &query);
 
         let target = TransformFrameIdHash::from_entity_path(&EntityPath::from("box"));
         let sources = [TransformFrameIdHash::from_entity_path(&EntityPath::from(