Move DataQueryResults into the ViewerContext (#4310)

### What - Part of the #4308 stack: - THIS PR: #4310 - #4311 - #4380 - #4381 Rather than doing queries on-demand, we now do all the queries up-front when creating the ViewerContext. This keeps us from re-running duplicate queries, and also eliminates the need to maintain a parallel implementation of resolve instead of just looking up the query results. Probably best reviewed commit-by-commit. This already hits some snags related to ViewerContext lifecycle that will hopefully be addressed in #1325 ### Checklist * [x] I have read and agree to [Contributor Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and the [Code of Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md) * [x] I've included a screenshot or gif (if applicable) * [x] I have tested [demo.rerun.io](https://demo.rerun.io/pr/4310) (if applicable) * [x] The PR title and labels are set such as to maximize their usefulness for the next release's CHANGELOG - [PR Build Summary](https://build.rerun.io/pr/4310) - [Docs preview](https://rerun.io/preview/2239f89de45fe0ac81fc5660ca13f3b95d7a8799/docs)  - [Examples preview](https://rerun.io/preview/2239f89de45fe0ac81fc5660ca13f3b95d7a8799/examples)  - [Recent benchmark results](https://build.rerun.io/graphs/crates.html) - [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)
rerun-io · Nov 30, 2023 · 0ffa61c · 0ffa61c
1 parent 3b88768
commit 0ffa61c
Show file tree

Hide file tree

Showing 15 changed files with 381 additions and 327 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/re_space_view/src/data_query.rs b/crates/re_space_view/src/data_query.rs
@@ -1,58 +1,5 @@
-use re_data_store::{EntityPath, EntityProperties, EntityPropertyMap};
-use re_viewer_context::{DataResult, EntitiesPerSystemPerClass, StoreContext};
-use slotmap::SlotMap;
-use smallvec::SmallVec;
-
-slotmap::new_key_type! {
-    /// Identifier for a [`DataResultNode`]
-    pub struct DataResultHandle;
-}
-
-/// A hierarchical tree of [`DataResult`]s
-pub struct DataResultTree {
-    pub data_results: SlotMap<DataResultHandle, DataResultNode>,
-    pub root_handle: Option<DataResultHandle>,
-}
-
-impl DataResultTree {
-    /// Depth-first traversal of the tree, calling `visitor` on each result.
-    pub fn visit(&self, visitor: &mut impl FnMut(DataResultHandle)) {
-        if let Some(root_handle) = self.root_handle {
-            self.visit_recursive(root_handle, visitor);
-        }
-    }
-
-    /// Look up a [`DataResult`] in the tree based on its handle.
-    pub fn lookup_result(&self, handle: DataResultHandle) -> Option<&DataResult> {
-        self.data_results.get(handle).map(|node| &node.data_result)
-    }
-
-    /// Look up a [`DataResultNode`] in the tree based on its handle.
-    pub fn lookup_node(&self, handle: DataResultHandle) -> Option<&DataResultNode> {
-        self.data_results.get(handle)
-    }
-
-    fn visit_recursive(
-        &self,
-        handle: DataResultHandle,
-        visitor: &mut impl FnMut(DataResultHandle),
-    ) {
-        if let Some(result) = self.data_results.get(handle) {
-            visitor(handle);
-
-            for child in &result.children {
-                self.visit_recursive(*child, visitor);
-            }
-        }
-    }
-}
-
-/// A single node in the [`DataResultTree`]
-#[derive(Debug)]
-pub struct DataResultNode {
-    pub data_result: DataResult,
-    pub children: SmallVec<[DataResultHandle; 4]>,
-}
+use re_data_store::{EntityProperties, EntityPropertyMap};
+use re_viewer_context::{DataQueryResult, EntitiesPerSystemPerClass, StoreContext};
 
 pub struct EntityOverrides {
     pub root: EntityProperties,
@@ -70,8 +17,8 @@ pub trait PropertyResolver {
 
 /// The common trait implemented for data queries
 ///
-/// Both interfaces return [`DataResult`]s, which are self-contained description of the data
-/// to be added to a `SpaceView` including both the [`EntityPath`] and context for any overrides.
+/// Both interfaces return [`re_viewer_context::DataResult`]s, which are self-contained description of the data
+/// to be added to a `SpaceView` including both the [`re_log_types::EntityPath`] and context for any overrides.
 pub trait DataQuery {
     /// Execute a full query, returning a `DataResultTree` containing all results.
     ///
@@ -83,20 +30,5 @@ pub trait DataQuery {
         property_resolver: &impl PropertyResolver,
         ctx: &StoreContext<'_>,
         entities_per_system_per_class: &EntitiesPerSystemPerClass,
-    ) -> DataResultTree;
-
-    /// Find a single [`DataResult`] within the context of the query.
-    ///
-    /// `auto_properties` is a map containing any heuristic-derived auto properties for the given `SpaceView`.
-    ///
-    /// This is used when finding the result for a single entity such as in
-    /// a selection panel.
-    fn resolve(
-        &self,
-        property_resolver: &impl PropertyResolver,
-        ctx: &StoreContext<'_>,
-        entities_per_system_per_class: &EntitiesPerSystemPerClass,
-        entity_path: &EntityPath,
-        as_group: bool,
-    ) -> DataResult;
+    ) -> DataQueryResult;
 }
diff --git a/crates/re_space_view/src/data_query_blueprint.rs b/crates/re_space_view/src/data_query_blueprint.rs
@@ -3,15 +3,13 @@ use once_cell::sync::Lazy;
 use re_data_store::{EntityProperties, EntityTree};
 use re_log_types::{EntityPath, EntityPathExpr};
 use re_viewer_context::{
-    DataResult, EntitiesPerSystem, EntitiesPerSystemPerClass, SpaceViewClassName,
+    DataQueryId, DataQueryResult, DataResult, DataResultHandle, DataResultNode, DataResultTree,
+    EntitiesPerSystem, EntitiesPerSystemPerClass, SpaceViewClassName,
 };
 use slotmap::SlotMap;
 use smallvec::SmallVec;
 
-use crate::{
-    blueprint::QueryExpressions, DataQuery, DataResultHandle, DataResultNode, DataResultTree,
-    EntityOverrides, PropertyResolver,
-};
+use crate::{blueprint::QueryExpressions, DataQuery, EntityOverrides, PropertyResolver};
 
 /// An implementation of [`DataQuery`] that is built from a collection of [`QueryExpressions`]
 ///
@@ -26,7 +24,7 @@ use crate::{
 /// and for which there is a valid `ViewPart` system. This keeps recursive expressions from incorrectly
 /// picking up irrelevant data within the tree.
 pub struct DataQueryBlueprint {
-    pub blueprint_path: EntityPath,
+    pub id: DataQueryId,
     pub space_view_class_name: SpaceViewClassName,
     pub expressions: QueryExpressions,
 }
@@ -41,7 +39,7 @@ impl DataQuery for DataQueryBlueprint {
         property_resolver: &impl PropertyResolver,
         ctx: &re_viewer_context::StoreContext<'_>,
         entities_per_system_per_class: &EntitiesPerSystemPerClass,
-    ) -> DataResultTree {
+    ) -> DataQueryResult {
         re_tracing::profile_function!();
 
         static EMPTY_ENTITY_LIST: Lazy<EntitiesPerSystem> = Lazy::new(Default::default);
@@ -66,56 +64,9 @@ impl DataQuery for DataQueryBlueprint {
             )
         });
 
-        DataResultTree {
-            data_results,
-            root_handle,
-        }
-    }
-
-    fn resolve(
-        &self,
-        property_resolver: &impl PropertyResolver,
-        ctx: &re_viewer_context::StoreContext<'_>,
-        entities_per_system_per_class: &EntitiesPerSystemPerClass,
-        entity_path: &re_log_types::EntityPath,
-        as_group: bool,
-    ) -> re_viewer_context::DataResult {
-        re_tracing::profile_function!();
-        let overrides = property_resolver.resolve_entity_overrides(ctx);
-
-        let view_parts = if let Some(per_system_entity_list) =
-            entities_per_system_per_class.get(&self.space_view_class_name)
-        {
-            per_system_entity_list
-                .iter()
-                .filter_map(|(part, ents)| {
-                    if ents.contains(entity_path) {
-                        Some(*part)
-                    } else {
-                        None
-                    }
-                })
-                .collect()
-        } else {
-            Default::default()
-        };
-
-        let mut resolved_properties = overrides.root.clone();
-        for prefix in EntityPath::incremental_walk(None, entity_path) {
-            resolved_properties = resolved_properties.with_child(&overrides.group.get(&prefix));
-        }
-
-        // TODO(jleibs): This needs to be updated to accommodate for groups
-        DataResult {
-            entity_path: entity_path.clone(),
-            view_parts,
-            is_group: as_group,
-            individual_properties: overrides.individual.get_opt(entity_path).cloned(),
-            resolved_properties,
-            override_path: self
-                .blueprint_path
-                .join(&Self::OVERRIDES_PREFIX.into())
-                .join(entity_path),
+        DataQueryResult {
+            id: self.id,
+            tree: DataResultTree::new(data_results, root_handle),
         }
     }
 }
@@ -214,7 +165,7 @@ impl<'a> QueryExpressionEvaluator<'a> {
             resolved_properties = resolved_properties.with_child(props);
         }
 
-        let base_entity_path = self.blueprint.blueprint_path.clone();
+        let base_entity_path = self.blueprint.id.as_entity_path().clone();
         let prefix = EntityPath::from(DataQueryBlueprint::OVERRIDES_PREFIX);
         let override_path = base_entity_path.join(&prefix).join(&entity_path);
 
@@ -387,7 +338,7 @@ mod tests {
 
         for (input, outputs) in scenarios {
             let query = DataQueryBlueprint {
-                blueprint_path: EntityPath::root(),
+                id: DataQueryId::random(),
                 space_view_class_name: "3D".into(),
                 expressions: input
                     .into_iter()
@@ -396,11 +347,11 @@ mod tests {
                     .into(),
             };
 
-            let result_tree = query.execute_query(&resolver, &ctx, &entities_per_system_per_class);
+            let query_result = query.execute_query(&resolver, &ctx, &entities_per_system_per_class);
 
             let mut visited = vec![];
-            result_tree.visit(&mut |handle| {
-                let result = result_tree.lookup_result(handle).unwrap();
+            query_result.tree.visit(&mut |handle| {
+                let result = query_result.tree.lookup_result(handle).unwrap();
                 if result.is_group && result.entity_path != EntityPath::root() {
                     visited.push(format!("{}/", result.entity_path));
                 } else {

diff --git a/crates/re_space_view/src/lib.rs b/crates/re_space_view/src/lib.rs
@@ -11,9 +11,7 @@ mod space_view_contents;
 mod unreachable_transform_reason;
 
 pub use blueprint::QueryExpressions;
-pub use data_query::{
-    DataQuery, DataResultHandle, DataResultNode, DataResultTree, EntityOverrides, PropertyResolver,
-};
+pub use data_query::{DataQuery, EntityOverrides, PropertyResolver};
 pub use data_query_blueprint::DataQueryBlueprint;
 pub use screenshot::ScreenshotMode;
 pub use space_view_contents::{DataBlueprintGroup, SpaceViewContents};

diff --git a/crates/re_space_view/src/space_view_contents.rs b/crates/re_space_view/src/space_view_contents.rs
@@ -3,15 +3,14 @@ use std::collections::{BTreeMap, BTreeSet};
 use nohash_hasher::IntMap;
 use re_data_store::{EntityPath, EntityProperties};
 use re_viewer_context::{
-    DataBlueprintGroupHandle, DataResult, EntitiesPerSystemPerClass, PerSystemEntities,
-    SpaceViewId, StoreContext, ViewSystemName,
+    DataBlueprintGroupHandle, DataQueryResult, DataResult, DataResultHandle, DataResultNode,
+    DataResultTree, EntitiesPerSystemPerClass, PerSystemEntities, SpaceViewId, StoreContext,
+    ViewSystemName,
 };
 use slotmap::SlotMap;
 use smallvec::{smallvec, SmallVec};
 
-use crate::{
-    DataQuery, DataResultHandle, DataResultNode, DataResultTree, EntityOverrides, PropertyResolver,
-};
+use crate::{DataQuery, EntityOverrides, PropertyResolver};
 
 /// A grouping of several data-blueprints.
 #[derive(Clone, serde::Deserialize, serde::Serialize)]
@@ -484,7 +483,7 @@ impl DataQuery for SpaceViewContents {
         property_resolver: &impl PropertyResolver,
         ctx: &StoreContext<'_>,
         _entities_per_system_per_class: &EntitiesPerSystemPerClass,
-    ) -> DataResultTree {
+    ) -> DataQueryResult {
         re_tracing::profile_function!();
         let overrides = property_resolver.resolve_entity_overrides(ctx);
         let mut data_results = SlotMap::<DataResultHandle, DataResultNode>::default();
@@ -495,74 +494,11 @@ impl DataQuery for SpaceViewContents {
             &overrides.root,
             &mut data_results,
         ));
-        DataResultTree {
-            data_results,
-            root_handle,
-        }
-    }
-
-    fn resolve(
-        &self,
-        property_resolver: &impl PropertyResolver,
-        ctx: &StoreContext<'_>,
-        _entities_per_system_per_class: &EntitiesPerSystemPerClass,
-        entity_path: &EntityPath,
-        as_group: bool,
-    ) -> DataResult {
-        re_tracing::profile_function!();
-        let overrides = property_resolver.resolve_entity_overrides(ctx);
-
-        let view_parts = self
-            .per_system_entities()
-            .iter()
-            .filter_map(|(part, ents)| {
-                if ents.contains(entity_path) {
-                    Some(*part)
-                } else {
-                    None
-                }
-            })
-            .collect();
-
-        // Start with the root override
-        let mut resolved_properties = overrides.root;
-
-        // Merge in any group overrides
-        for prefix in EntityPath::incremental_walk(None, entity_path) {
-            if let Some(props) = overrides.group.get_opt(&prefix) {
-                resolved_properties = resolved_properties.with_child(props);
-            }
-        }
-
-        if as_group {
-            DataResult {
-                entity_path: entity_path.clone(),
-                view_parts,
-                is_group: true,
-                resolved_properties,
-                individual_properties: overrides.group.get_opt(entity_path).cloned(),
-                override_path: self
-                    .entity_path()
-                    .join(&SpaceViewContents::GROUP_OVERRIDES_PREFIX.into())
-                    .join(entity_path),
-            }
-        } else {
-            // Finally apply the individual overrides
-            if let Some(props) = overrides.individual.get_opt(entity_path) {
-                resolved_properties = resolved_properties.with_child(props);
-            }
-
-            DataResult {
-                entity_path: entity_path.clone(),
-                view_parts,
-                is_group: false,
-                resolved_properties,
-                individual_properties: overrides.individual.get_opt(entity_path).cloned(),
-                override_path: self
-                    .entity_path()
-                    .join(&SpaceViewContents::INDIVIDUAL_OVERRIDES_PREFIX.into())
-                    .join(entity_path),
-            }
+        DataQueryResult {
+            // Create a fake `DataQueryId` based on the SpaceView since each
+            // SpaceView contains a single "Query" based on its contents.
+            id: self.space_view_id.uuid().into(),
+            tree: DataResultTree::new(data_results, root_handle),
         }
     }
 }

diff --git a/crates/re_viewer/Cargo.toml b/crates/re_viewer/Cargo.toml
@@ -51,6 +51,7 @@ re_log.workspace = true
 re_memory.workspace = true
 re_renderer = { workspace = true, default-features = false }
 re_smart_channel.workspace = true
+re_space_view.workspace = true
 re_space_view_bar_chart.workspace = true
 re_space_view_spatial.workspace = true
 re_space_view_tensor.workspace = true
@@ -87,6 +88,7 @@ egui.workspace = true
 image = { workspace = true, default-features = false, features = ["png"] }
 itertools = { workspace = true }
 once_cell = { workspace = true }
+nohash-hasher.workspace = true
 poll-promise = { workspace = true, features = ["web"] }
 rfd.workspace = true
 ron.workspace = true