Improve performance for many entities

crates/re_log_types/src/serde_field.rs

@@ -63,14 +63,17 @@ where
 
 impl<T> ArrowDeserialize for SerdeField<T>
 where
-    T: serde::ser::Serialize + serde::de::DeserializeOwned,
+    T: serde::ser::Serialize + serde::de::DeserializeOwned + Default,
 {
     type ArrayType = BinaryArray<i32>;
 
     #[inline]
     fn arrow_deserialize(v: <&Self::ArrayType as IntoIterator>::Item) -> Option<T> {
         re_tracing::profile_function!();
-        v.and_then(|v| rmp_serde::from_slice::<T>(v).ok())
+        Some(
+            v.and_then(|v| rmp_serde::from_slice::<T>(v).ok())
+                .unwrap_or_default(),
+        )
     }
 }
 
@@ -79,7 +82,7 @@ mod tests {
     use super::SerdeField;
     use arrow2_convert::{ArrowDeserialize, ArrowField, ArrowSerialize};
 
-    #[derive(Clone, Debug, PartialEq, serde::Deserialize, serde::Serialize)]
+    #[derive(Clone, Debug, PartialEq, Default, serde::Deserialize, serde::Serialize)]
     struct SomeStruct {
         foo: String,
         bar: u32,

crates/re_space_view/src/lib.rs

@@ -3,10 +3,10 @@
 //! Types & utilities for defining Space View classes and communicating with the Viewport.
 
 pub mod controls;
-mod data_blueprint;
 mod screenshot;
+mod space_view_contents;
 mod unreachable_transform_reason;
 
-pub use data_blueprint::{DataBlueprintGroup, DataBlueprintTree};
 pub use screenshot::ScreenshotMode;
+pub use space_view_contents::{DataBlueprintGroup, SpaceViewContents};
 pub use unreachable_transform_reason::UnreachableTransformReason;

crates/re_space_view/src/data_blueprint.rs → crates/re_space_view/src/space_view_contents.rs

@@ -1,8 +1,8 @@
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 
-use nohash_hasher::{IntMap, IntSet};
+use nohash_hasher::IntMap;
 use re_data_store::{EntityPath, EntityProperties, EntityPropertyMap};
-use re_viewer_context::DataBlueprintGroupHandle;
+use re_viewer_context::{DataBlueprintGroupHandle, PerSystemEntities, ViewSystemName};
 use slotmap::SlotMap;
 use smallvec::{smallvec, SmallVec};
 
@@ -79,7 +79,7 @@ impl DataBlueprints {
 
 /// Tree of all data blueprint groups for a single space view.
 #[derive(Clone, serde::Deserialize, serde::Serialize)]
-pub struct DataBlueprintTree {
+pub struct SpaceViewContents {
     /// All data blueprint groups.
     groups: SlotMap<DataBlueprintGroupHandle, DataBlueprintGroup>,
 
@@ -95,8 +95,7 @@ pub struct DataBlueprintTree {
     /// Two things to keep in sync:
     /// * children on [`DataBlueprintGroup`] this is on
     /// * elements in [`Self::path_to_group`]
-    /// TODO(andreas): Can we reduce the amount of these dependencies?
-    entity_paths: IntSet<EntityPath>,
+    per_system_entity_list: PerSystemEntities,
 
     /// Root group, always exists as a placeholder
     root_group_handle: DataBlueprintGroupHandle,
@@ -105,12 +104,12 @@ pub struct DataBlueprintTree {
 }
 
 /// Determine whether this `DataBlueprintTree` has user-edits relative to another `DataBlueprintTree`
-impl DataBlueprintTree {
+impl SpaceViewContents {
     pub fn has_edits(&self, other: &Self) -> bool {
         let Self {
             groups,
             path_to_group,
-            entity_paths,
+            per_system_entity_list: entity_paths,
             root_group_handle,
             data_blueprints,
         } = self;
@@ -123,13 +122,13 @@ impl DataBlueprintTree {
                     .map_or(true, |other_val| val.has_edits(other_val))
             })
             || *path_to_group != other.path_to_group
-            || *entity_paths != other.entity_paths
+            || *entity_paths != other.per_system_entity_list
             || *root_group_handle != other.root_group_handle
             || data_blueprints.has_edits(&other.data_blueprints)
     }
 }
 
-impl Default for DataBlueprintTree {
+impl Default for SpaceViewContents {
     fn default() -> Self {
         let mut groups = SlotMap::default();
         let root_group = groups.insert(DataBlueprintGroup::default());
@@ -140,14 +139,14 @@ impl Default for DataBlueprintTree {
         Self {
             groups,
             path_to_group: path_to_blueprint,
-            entity_paths: IntSet::default(),
+            per_system_entity_list: BTreeMap::default(),
             root_group_handle: root_group,
             data_blueprints: DataBlueprints::default(),
         }
     }
 }
 
-impl DataBlueprintTree {
+impl SpaceViewContents {
     /// Returns a handle to the root data blueprint.
     ///
     /// Even if there are no other groups, we always have a root group at the top.
@@ -207,12 +206,47 @@ impl DataBlueprintTree {
     }
 
     pub fn contains_entity(&self, path: &EntityPath) -> bool {
-        self.entity_paths.contains(path)
+        // If an entity is in path_to_group it is *likely* also an entity in the Space View.
+        // However, it could be that the path *only* refers to a group, not also an entity.
+        // So once we resolved the group, we need to check if it contains the entity of interest.
+        self.path_to_group
+            .get(path)
+            .and_then(|group| {
+                self.groups
+                    .get(*group)
+                    .and_then(|group| group.entities.get(path))
+            })
+            .is_some()
     }
 
     /// List of all entities that we query via this data blueprint collection.
-    pub fn entity_paths(&self) -> &IntSet<EntityPath> {
-        &self.entity_paths
+    pub fn entity_paths(&self) -> impl Iterator<Item = &EntityPath> {
+        // Each entity is only ever in one group, therefore collecting all entities from all groups, gives us all entities.
+        self.groups.values().flat_map(|group| group.entities.iter())
+    }
+
+    pub fn per_system_entities(&self) -> &PerSystemEntities {
+        &self.per_system_entity_list
+    }
+
+    pub fn per_system_entities_mut(&mut self) -> &mut PerSystemEntities {
+        &mut self.per_system_entity_list
+    }
+
+    pub fn contains_all_entities_from(&self, other: &SpaceViewContents) -> bool {
+        for (system, entities) in &other.per_system_entity_list {
+            let Some(self_entities) = self.per_system_entity_list.get(system) else {
+                if entities.is_empty() {
+                    continue;
+                } else {
+                    return false;
+                }
+            };
+            if !entities.is_subset(self_entities) {
+                return false;
+            }
+        }
+        true
     }
 
     /// Should be called on frame start.
@@ -225,7 +259,7 @@ impl DataBlueprintTree {
         // and/or when new entity paths are added, but such memoization would add complexity.
 
         fn project_tree(
-            tree: &mut DataBlueprintTree,
+            tree: &mut SpaceViewContents,
             parent_properties: &EntityProperties,
             group_handle: DataBlueprintGroupHandle,
         ) {
@@ -270,7 +304,14 @@ impl DataBlueprintTree {
         let mut new_leaf_groups = Vec::new();
 
         for path in paths {
-            self.entity_paths.insert(path.clone());
+            // TODO(andreas/jleibs): Incoming entities should already "know" what system they belong to.
+            // WARNING: This is less clear with ContextSystems.
+            // Generally, we don't want the user to have control over context systems since they "prepare" data for PartSystems.
+            // -> Context systems should be a direct consequence of the PartSystems that are active for a given entity.
+            self.per_system_entity_list
+                .entry(ViewSystemName::unknown())
+                .or_default()
+                .insert(path.clone());
 
             // Is there already a group associated with this exact path? (maybe because a child was logged there earlier)
             // If so, we can simply move it under this existing group.
@@ -393,14 +434,19 @@ impl DataBlueprintTree {
     ///
     /// If the entity was not known by this data blueprint tree nothing happens.
     pub fn remove_entity(&mut self, path: &EntityPath) {
+        re_tracing::profile_function!();
+
         if let Some(group_handle) = self.path_to_group.get(path) {
             if let Some(group) = self.groups.get_mut(*group_handle) {
                 group.entities.remove(path);
                 self.remove_group_if_empty(*group_handle);
             }
         }
         self.path_to_group.remove(path);
-        self.entity_paths.remove(path);
+
+        for per_system_list in self.per_system_entity_list.values_mut() {
+            per_system_list.remove(path);
+        }
     }
 
     /// Removes a group and all its entities and subgroups from the blueprint tree
@@ -421,7 +467,9 @@ impl DataBlueprintTree {
 
         // Remove all child entities.
         for entity_path in &group.entities {
-            self.entity_paths.remove(entity_path);
+            for per_system_list in self.per_system_entity_list.values_mut() {
+                per_system_list.remove(entity_path);
+            }
         }
 
         // Remove from `path_to_group` map.

crates/re_space_view_bar_chart/src/view_part_system.rs

@@ -6,8 +6,8 @@ use re_data_store::EntityPath;
 use re_log_types::{TimeInt, Timeline};
 use re_types::{ComponentName, Loggable as _};
 use re_viewer_context::{
-    ArchetypeDefinition, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartSystem,
-    ViewQuery, ViewerContext,
+    ArchetypeDefinition, NamedViewSystem, SpaceViewSystemExecutionError, ViewContextCollection,
+    ViewPartSystem, ViewQuery, ViewerContext,
 };
 
 /// A bar chart system, with everything needed to render it.
@@ -16,6 +16,12 @@ pub struct BarChartViewPartSystem {
     pub charts: BTreeMap<EntityPath, Tensor>,
 }
 
+impl NamedViewSystem for BarChartViewPartSystem {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "BarChartView".into()
+    }
+}
+
 impl ViewPartSystem for BarChartViewPartSystem {
     fn archetype(&self) -> ArchetypeDefinition {
         vec1::vec1![Tensor::name()]
@@ -51,7 +57,7 @@ impl ViewPartSystem for BarChartViewPartSystem {
 
         let store = &ctx.store_db.entity_db.data_store;
 
-        for (ent_path, props) in query.iter_entities() {
+        for (ent_path, props) in query.iter_entities_for_system(Self::name()) {
             if !props.visible {
                 continue;
             }

crates/re_space_view_spatial/src/contexts/annotation_context.rs

@@ -1,9 +1,17 @@
 use re_types::{components::AnnotationContext, Loggable};
-use re_viewer_context::{AnnotationMap, ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{
+    AnnotationMap, ArchetypeDefinition, NamedViewSystem, ViewContextSystem, ViewSystemName,
+};
 
 #[derive(Default)]
 pub struct AnnotationSceneContext(pub AnnotationMap);
 
+impl NamedViewSystem for AnnotationSceneContext {
+    fn name() -> ViewSystemName {
+        "AnnotationSceneContext".into()
+    }
+}
+
 impl ViewContextSystem for AnnotationSceneContext {
     fn archetypes(&self) -> Vec<ArchetypeDefinition> {
         vec![vec1::vec1![AnnotationContext::name()]]
@@ -14,11 +22,15 @@ impl ViewContextSystem for AnnotationSceneContext {
         ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
-        self.0.load(
-            ctx,
-            &query.latest_at_query(),
-            query.iter_entities().map(|(p, _)| p),
-        );
+        // 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.
+        // TODO(andreas): This curiously seems to be the only place where we get a view query and are interested in _all_ entities.
+        //                We know that this deals with duplicates entities just fine, so we just concat all per-system entities here.
+        let entities = query
+            .per_system_entities
+            .values()
+            .flat_map(|entities| entities.iter());
+        self.0.load(ctx, &query.latest_at_query(), entities);
     }
 
     fn as_any(&self) -> &dyn std::any::Any {

crates/re_space_view_spatial/src/contexts/depth_offsets.rs

@@ -3,7 +3,7 @@ use std::collections::{BTreeMap, BTreeSet};
 use nohash_hasher::IntMap;
 use re_log_types::EntityPathHash;
 use re_types::{components::DrawOrder, Loggable as _};
-use re_viewer_context::{ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{ArchetypeDefinition, NamedViewSystem, ViewContextSystem};
 
 /// Context for creating a mapping from [`DrawOrder`] to [`re_renderer::DepthOffset`].
 #[derive(Default)]
@@ -18,6 +18,12 @@ pub struct EntityDepthOffsets {
     pub points: re_renderer::DepthOffset,
 }
 
+impl NamedViewSystem for EntityDepthOffsets {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "EntityDepthOffsets".into()
+    }
+}
+
 impl ViewContextSystem for EntityDepthOffsets {
     fn archetypes(&self) -> Vec<ArchetypeDefinition> {
         vec![vec1::vec1![DrawOrder::name()]]
@@ -28,8 +34,6 @@ impl ViewContextSystem for EntityDepthOffsets {
         ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
-        re_tracing::profile_function!();
-
         #[derive(PartialEq, PartialOrd, Eq, Ord)]
         enum DrawOrderTarget {
             Entity(EntityPathHash),
@@ -43,7 +47,7 @@ impl ViewContextSystem for EntityDepthOffsets {
 
         // Use a BTreeSet for entity hashes to get a stable order.
         let mut entities_per_draw_order = BTreeMap::<DrawOrder, BTreeSet<DrawOrderTarget>>::new();
-        for (ent_path, _) in query.iter_entities() {
+        for (ent_path, _) in query.iter_entities_for_system(Self::name()) {
             if let Some(draw_order) = store.query_latest_component::<DrawOrder>(
                 ent_path,
                 &ctx.rec_cfg.time_ctrl.current_query(),

crates/re_space_view_spatial/src/contexts/mod.rs

@@ -15,7 +15,9 @@ pub use transform_context::{pinhole_camera_view_coordinates, TransformContext};
 // -----------------------------------------------------------------------------
 
 use re_renderer::DepthOffset;
-use re_viewer_context::{Annotations, SpaceViewClassRegistryError, ViewContextSystem};
+use re_viewer_context::{
+    Annotations, NamedViewSystem, SpaceViewClassRegistryError, ViewContextSystem,
+};
 
 /// Context objects for a single entity in a spatial scene.
 pub struct SpatialSceneEntityContext<'a> {
@@ -32,6 +34,12 @@ pub struct PrimitiveCounter {
     pub num_primitives: AtomicUsize,
 }
 
+impl NamedViewSystem for PrimitiveCounter {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "PrimitiveCounter".into()
+    }
+}
+
 impl ViewContextSystem for PrimitiveCounter {
     fn archetypes(&self) -> Vec<re_viewer_context::ArchetypeDefinition> {
         Vec::new()