Improve many-entity performance by fixing redundant AnnotationContext lookups (also affects scene without annotation contexts) (#11445)

Author: Wumpf

Cargo.lock

@@ -10170,6 +10170,7 @@ dependencies = [
  "ndarray",
  "nohash-hasher",
  "parking_lot",
+ "rayon",
  "re_arrow_ui",
  "re_arrow_util",
  "re_byte_size",

crates/viewer/re_data_ui/src/lib.rs

@@ -148,7 +148,7 @@ pub fn annotations(
 ) -> std::sync::Arc<re_viewer_context::Annotations> {
     re_tracing::profile_function!();
     let mut annotation_map = re_viewer_context::AnnotationMap::default();
-    annotation_map.load(ctx, query, std::iter::once(entity_path));
+    annotation_map.load(ctx, query);
     annotation_map.find(entity_path)
 }
 

crates/viewer/re_test_viewport/src/lib.rs

@@ -151,15 +151,21 @@ impl TestContextExt for TestContext {
             ViewBlueprint::try_from_db(view_id, ctx.store_context.blueprint, ctx.blueprint_query)
                 .expect("expected the view id to be known to the blueprint store");
 
-        let view_class = ctx
-            .view_class_registry()
-            .get_class_or_log_error(view_blueprint.class_identifier());
+        let class_registry = ctx.view_class_registry();
+        let class_identifier = view_blueprint.class_identifier();
+        let view_class = class_registry.get_class_or_log_error(class_identifier);
 
         let mut view_states = self.view_states.lock();
         let view_state = view_states.get_mut_or_create(view_id, view_class);
 
-        let (view_query, system_execution_output) =
-            execute_systems_for_view(ctx, &view_blueprint, view_state);
+        let context_system_static_exec_results = class_registry
+            .run_static_context_systems_for_views(ctx, std::iter::once(class_identifier));
+        let (view_query, system_execution_output) = execute_systems_for_view(
+            ctx,
+            &view_blueprint,
+            view_state,
+            &context_system_static_exec_results,
+        );
 
         view_class
             .ui(ctx, ui, view_state, &view_query, system_execution_output)

crates/viewer/re_view/src/annotation_scene_context.rs

@@ -1,9 +1,12 @@
+use std::sync::Arc;
+
 use re_viewer_context::{
-    AnnotationMap, IdentifiedViewSystem, ViewContextSystem, ViewSystemIdentifier,
+    AnnotationMap, IdentifiedViewSystem, ViewContextSystem, ViewContextSystemStaticExecResult,
+    ViewSystemIdentifier,
 };
 
 #[derive(Default)]
-pub struct AnnotationSceneContext(pub AnnotationMap);
+pub struct AnnotationSceneContext(pub Arc<AnnotationMap>);
 
 impl IdentifiedViewSystem for AnnotationSceneContext {
     fn identifier() -> ViewSystemIdentifier {
@@ -12,19 +15,30 @@ impl IdentifiedViewSystem for AnnotationSceneContext {
 }
 
 impl ViewContextSystem for AnnotationSceneContext {
+    fn execute_static(
+        ctx: &re_viewer_context::ViewerContext<'_>,
+    ) -> ViewContextSystemStaticExecResult {
+        // Use static execution to load the annotation map for all entities.
+        // Alternatively, we could do this only for visible ones per View but this is actually a lot more expensive to do
+        // given that there's typically just one annotation map per recording anyways!
+        let mut annotation_map = AnnotationMap::default();
+        annotation_map.load(ctx, &ctx.current_query());
+
+        Box::new(Self(Arc::new(annotation_map)))
+    }
+
     fn execute(
         &mut self,
-        ctx: &re_viewer_context::ViewContext<'_>,
-        query: &re_viewer_context::ViewQuery<'_>,
+        _ctx: &re_viewer_context::ViewContext<'_>,
+        _query: &re_viewer_context::ViewQuery<'_>,
+        static_execution_result: &ViewContextSystemStaticExecResult,
     ) {
-        re_tracing::profile_function!();
-        // We create a list of *all* entities here, do not only iterate over those with annotation context.
-        // TODO(andreas): But knowing ahead of time where we have annotation contexts could be used for optimization.
-        self.0.load(
-            ctx.viewer_ctx,
-            &query.latest_at_query(),
-            query.iter_all_entities(),
-        );
+        // Take over the static result to make it available.
+        self.0 = static_execution_result
+            .downcast_ref::<Self>()
+            .expect("Unexpected static execution result type")
+            .0
+            .clone();
     }
 
     fn as_any(&self) -> &dyn std::any::Any {

crates/viewer/re_view_map/src/map_view.rs

@@ -122,6 +122,7 @@ impl ViewClass for MapView {
         system_registry.register_visualizer::<GeoLineStringsVisualizer>()?;
 
         system_registry.register_context_system::<AnnotationSceneContext>()?;
+        re_viewer_context::AnnotationContextStoreSubscriber::subscription_handle(); // Needed by `AnnotationSceneContext`
 
         Ok(())
     }

crates/viewer/re_view_spatial/src/contexts/depth_offsets.rs

@@ -29,6 +29,7 @@ impl ViewContextSystem for EntityDepthOffsets {
         &mut self,
         ctx: &re_viewer_context::ViewContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
+        _static_execution_result: &re_viewer_context::ViewContextSystemStaticExecResult,
     ) {
         let mut entities_per_draw_order = BTreeMap::new();
         for (visualizer, draw_order_descriptor) in visualizers_processing_draw_order() {

crates/viewer/re_view_spatial/src/contexts/mod.rs

@@ -28,6 +28,9 @@ pub fn register_spatial_contexts(
 ) -> Result<(), ViewClassRegistryError> {
     system_registry.register_context_system::<TransformTreeContext>()?;
     system_registry.register_context_system::<EntityDepthOffsets>()?;
+
     system_registry.register_context_system::<AnnotationSceneContext>()?;
+    re_viewer_context::AnnotationContextStoreSubscriber::subscription_handle(); // Needed by `AnnotationSceneContext`
+
     Ok(())
 }

crates/viewer/re_view_spatial/src/contexts/transform_tree_context.rs

@@ -2,7 +2,10 @@ use re_chunk_store::LatestAtQuery;
 use re_log_types::{EntityPath, EntityPathHash};
 use re_types::{archetypes, components::ImagePlaneDistance};
 use re_view::DataResultQuery as _;
-use re_viewer_context::{DataResultTree, IdentifiedViewSystem, ViewContext, ViewContextSystem};
+use re_viewer_context::{
+    DataResultTree, IdentifiedViewSystem, ViewContext, ViewContextSystem,
+    ViewContextSystemStaticExecResult,
+};
 
 use crate::{caches::TransformDatabaseStoreCache, visualizers::CamerasVisualizer};
 
@@ -20,6 +23,7 @@ impl ViewContextSystem for TransformTreeContext {
         &mut self,
         ctx: &re_viewer_context::ViewContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
+        _static_execution_result: &ViewContextSystemStaticExecResult,
     ) {
         let recording = ctx.recording();
 

crates/viewer/re_view_spatial/src/max_image_dimension_subscriber.rs

@@ -55,9 +55,7 @@ impl MaxImageDimensionsStoreSubscriber {
             move |subscriber: &Self| f(&subscriber.max_dimensions),
         )
     }
-}
 
-impl MaxImageDimensionsStoreSubscriber {
     /// Accesses the global store subscriber.
     ///
     /// Lazily registers the subscriber if it hasn't been registered yet.

crates/viewer/re_viewer_context/Cargo.toml

@@ -68,6 +68,7 @@ macaw.workspace = true
 ndarray.workspace = true
 nohash-hasher.workspace = true
 parking_lot = { workspace = true, features = ["serde"] }
+rayon.workspace = true
 serde.workspace = true
 slotmap.workspace = true
 smallvec.workspace = true