add schema_for_view_contents (antoine happy)

rerun-io · Oct 2, 2024 · ab06ff8 · ab06ff8
1 parent 9412150
commit ab06ff8
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 19 deletions.
diff --git a/crates/store/re_chunk_store/src/dataframe.rs b/crates/store/re_chunk_store/src/dataframe.rs
@@ -748,7 +748,7 @@ impl std::fmt::Display for SparseFillStrategy {
 ///
 // TODO(cmc): we need to be able to build that easily in a command-line context, otherwise it's just
 // very annoying. E.g. `--with /world/points:[rr.Position3D, rr.Radius] --with /cam:[rr.Pinhole]`.
-pub type ViewContents = BTreeMap<EntityPath, Option<BTreeSet<ComponentName>>>;
+pub type ViewContentsSelector = BTreeMap<EntityPath, Option<BTreeSet<ComponentName>>>;
 
 // TODO(cmc): Ultimately, this shouldn't be hardcoded to `Timeline`, but to a generic `I: Index`.
 //            `Index` in this case should also be implemented on tuples (`(I1, I2, ...)`).
@@ -801,7 +801,7 @@ pub struct QueryExpression2 {
     ///   "metrics": [rr.Scalar]
     /// }
     /// ```
-    pub view_contents: Option<ViewContents>,
+    pub view_contents: Option<ViewContentsSelector>,
 
     /// The index used to filter out _rows_ from the view contents.
     ///
@@ -1153,4 +1153,31 @@ impl ChunkStore {
             .filter(|descr| !filtered_out.contains(descr))
             .collect()
     }
+
+    /// Returns the filtered schema for the given [`ViewContentsSelector`].
+    ///
+    /// The order of the columns is guaranteed to be in a specific order:
+    /// * first, the control columns in lexical order (`RowId`);
+    /// * second, the time columns in lexical order (`frame_nr`, `log_time`, ...);
+    /// * third, the component columns in lexical order (`Color`, `Radius, ...`).
+    pub fn schema_for_view_contents(
+        &self,
+        view_contents: &ViewContentsSelector,
+    ) -> Vec<ColumnDescriptor> {
+        re_tracing::profile_function!();
+
+        self.schema()
+            .into_iter()
+            .filter(|column| match column {
+                ColumnDescriptor::Control(_) | ColumnDescriptor::Time(_) => true,
+                ColumnDescriptor::Component(column) => view_contents
+                    .get(&column.entity_path)
+                    .map_or(false, |components| {
+                        components.as_ref().map_or(true, |components| {
+                            components.contains(&column.component_name)
+                        })
+                    }),
+            })
+            .collect()
+    }
 }
diff --git a/crates/store/re_chunk_store/src/lib.rs b/crates/store/re_chunk_store/src/lib.rs
@@ -27,7 +27,7 @@ pub use self::dataframe::{
     ColumnDescriptor, ColumnSelector, ComponentColumnDescriptor, ComponentColumnSelector,
     ControlColumnDescriptor, ControlColumnSelector, Index, IndexRange, IndexValue, JoinEncoding,
     LatestAtQueryExpression, QueryExpression, QueryExpression2, RangeQueryExpression,
-    SparseFillStrategy, TimeColumnDescriptor, TimeColumnSelector,
+    SparseFillStrategy, TimeColumnDescriptor, TimeColumnSelector, ViewContentsSelector,
 };
 pub use self::events::{ChunkStoreDiff, ChunkStoreDiffKind, ChunkStoreEvent};
 pub use self::gc::{GarbageCollectionOptions, GarbageCollectionTarget};

diff --git a/crates/store/re_dataframe2/src/engine.rs b/crates/store/re_dataframe2/src/engine.rs
@@ -1,5 +1,7 @@
 use re_chunk::{EntityPath, TransportChunk};
-use re_chunk_store::{ChunkStore, ColumnDescriptor, QueryExpression, QueryExpression2};
+use re_chunk_store::{
+    ChunkStore, ColumnDescriptor, QueryExpression, QueryExpression2, ViewContentsSelector,
+};
 use re_log_types::EntityPathFilter;
 use re_query::Caches;
 
@@ -66,6 +68,14 @@ impl QueryEngine<'_> {
         self.store.schema_for_query(query)
     }
 
+    #[inline]
+    pub fn schema_for_view_contents(
+        &self,
+        view_contents: &ViewContentsSelector,
+    ) -> Vec<ColumnDescriptor> {
+        self.store.schema_for_view_contents(view_contents)
+    }
+
     /// Starts a new query by instantiating a [`QueryHandle`].
     #[inline]
     pub fn query(&self, query: QueryExpression2) -> QueryHandle<'_> {

diff --git a/crates/store/re_dataframe2/src/query.rs b/crates/store/re_dataframe2/src/query.rs
@@ -125,21 +125,7 @@ impl QueryHandle<'_> {
 
             // 1. Compute the schema of the view contents.
             let view_contents = if let Some(view_contents) = self.query.view_contents.as_ref() {
-                self.engine
-                    .store
-                    .schema()
-                    .into_iter()
-                    .filter(|column| match column {
-                        ColumnDescriptor::Control(_) | ColumnDescriptor::Time(_) => true,
-                        ColumnDescriptor::Component(column) => view_contents
-                            .get(&column.entity_path)
-                            .map_or(false, |components| {
-                                components.as_ref().map_or(true, |components| {
-                                    components.contains(&column.component_name)
-                                })
-                            }),
-                    })
-                    .collect_vec()
+                self.engine.store.schema_for_view_contents(view_contents)
             } else {
                 self.engine.store.schema()
             };