Allow systems to specify custom filter for the applicable set beyond required_components

crates/re_space_view_bar_chart/src/view_part_system.rs

@@ -6,11 +6,11 @@ use re_types::{
     archetypes::{BarChart, Tensor},
     components::Color,
     datatypes::TensorData,
-    Archetype, ComponentNameSet,
+    Archetype, ComponentNameSet, Loggable,
 };
 use re_viewer_context::{
-    default_heuristic_filter, HeuristicFilterContext, IdentifiedViewSystem,
-    SpaceViewSystemExecutionError, ViewContextCollection, ViewPartSystem, ViewQuery, ViewerContext,
+    IdentifiedViewSystem, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartSystem,
+    ViewQuery, ViewerContext, VisualizerAdditionalApplicabilityFilter,
 };
 
 /// A bar chart system, with everything needed to render it.
@@ -25,6 +25,19 @@ impl IdentifiedViewSystem for BarChartViewPartSystem {
     }
 }
 
+struct BarChartVisualizerEntityFilter;
+
+impl VisualizerAdditionalApplicabilityFilter for BarChartVisualizerEntityFilter {
+    fn update_applicability(&mut self, event: &re_arrow_store::StoreEvent) -> bool {
+        event.diff.cells.iter().any(|(component_name, cell)| {
+            component_name == &re_types::components::TensorData::name()
+                && re_types::components::TensorData::from_arrow(cell.as_arrow_ref())
+                    .map(|tensors| tensors.iter().any(|tensor| tensor.is_vector()))
+                    .unwrap_or(false)
+        })
+    }
+}
+
 impl ViewPartSystem for BarChartViewPartSystem {
     fn required_components(&self) -> ComponentNameSet {
         BarChart::required_components()
@@ -42,28 +55,8 @@ impl ViewPartSystem for BarChartViewPartSystem {
             .collect()
     }
 
-    fn heuristic_filter(
-        &self,
-        store: &re_arrow_store::DataStore,
-        ent_path: &EntityPath,
-        _ctx: HeuristicFilterContext,
-        query: &LatestAtQuery,
-        entity_components: &ComponentNameSet,
-    ) -> bool {
-        if !default_heuristic_filter(entity_components, &self.indicator_components()) {
-            return false;
-        }
-
-        // NOTE: We want to make sure we query at the right time, otherwise we might take into
-        // account a `Clear()` that actually only applies into the future, and then
-        // `is_vector` will righfully fail because of the empty tensor.
-        if let Some(tensor) =
-            store.query_latest_component::<re_types::components::TensorData>(ent_path, query)
-        {
-            tensor.is_vector()
-        } else {
-            false
-        }
+    fn applicability_filter(&self) -> Option<Box<dyn VisualizerAdditionalApplicabilityFilter>> {
+        Some(Box::new(BarChartVisualizerEntityFilter))
     }
 
     fn execute(

crates/re_space_view_spatial/src/parts/images.rs

@@ -16,12 +16,12 @@ use re_types::{
     archetypes::{DepthImage, Image, SegmentationImage},
     components::{Color, DrawOrder, TensorData, ViewCoordinates},
     tensor_data::{DecodedTensor, TensorDataMeaning},
-    Archetype as _, ComponentNameSet,
+    Archetype as _, ComponentNameSet, Loggable as _,
 };
 use re_viewer_context::{
     default_heuristic_filter, gpu_bridge, DefaultColor, HeuristicFilterContext, SpaceViewClass,
     SpaceViewSystemExecutionError, TensorDecodeCache, TensorStatsCache, ViewPartSystem, ViewQuery,
-    ViewerContext,
+    ViewerContext, VisualizerAdditionalApplicabilityFilter,
 };
 use re_viewer_context::{IdentifiedViewSystem, ViewContextCollection};
 
@@ -643,6 +643,23 @@ impl IdentifiedViewSystem for ImagesPart {
     }
 }
 
+struct ImageVisualizerEntityFilter;
+
+impl VisualizerAdditionalApplicabilityFilter for ImageVisualizerEntityFilter {
+    fn update_applicability(&mut self, event: &re_arrow_store::StoreEvent) -> bool {
+        event.diff.cells.iter().any(|(component_name, cell)| {
+            component_name == &re_types::components::TensorData::name()
+                && re_types::components::TensorData::from_arrow(cell.as_arrow_ref())
+                    .map(|tensors| {
+                        tensors
+                            .iter()
+                            .any(|tensor| tensor.is_shaped_like_an_image())
+                    })
+                    .unwrap_or(false)
+        })
+    }
+}
+
 impl ViewPartSystem for ImagesPart {
     fn required_components(&self) -> ComponentNameSet {
         let image: ComponentNameSet = Image::required_components()
@@ -677,12 +694,16 @@ impl ViewPartSystem for ImagesPart {
         .collect()
     }
 
+    fn applicability_filter(&self) -> Option<Box<dyn VisualizerAdditionalApplicabilityFilter>> {
+        Some(Box::new(ImageVisualizerEntityFilter))
+    }
+
     fn heuristic_filter(
         &self,
-        store: &re_arrow_store::DataStore,
-        ent_path: &EntityPath,
+        _store: &re_arrow_store::DataStore,
+        _ent_path: &EntityPath,
         ctx: HeuristicFilterContext,
-        query: &LatestAtQuery,
+        _query: &LatestAtQuery,
         entity_components: &ComponentNameSet,
     ) -> bool {
         if !default_heuristic_filter(entity_components, &self.indicator_components()) {
@@ -697,14 +718,7 @@ impl ViewPartSystem for ImagesPart {
             return false;
         }
 
-        // NOTE: We want to make sure we query at the right time, otherwise we might take into
-        // account a `Clear()` that actually only applies into the future, and then
-        // `is_shaped_like_an_image` will righfully fail because of the empty tensor.
-        if let Some(tensor) = store.query_latest_component::<TensorData>(ent_path, query) {
-            tensor.is_shaped_like_an_image()
-        } else {
-            false
-        }
+        true
     }
 
     fn execute(

crates/re_space_view_tensor/src/view_part_system.rs

@@ -3,12 +3,11 @@ use re_data_store::EntityPath;
 use re_log_types::RowId;
 use re_types::{
     archetypes::Tensor, components::TensorData, tensor_data::DecodedTensor, Archetype,
-    ComponentNameSet,
+    ComponentNameSet, Loggable as _,
 };
 use re_viewer_context::{
-    default_heuristic_filter, HeuristicFilterContext, IdentifiedViewSystem,
-    SpaceViewSystemExecutionError, TensorDecodeCache, ViewContextCollection, ViewPartSystem,
-    ViewQuery, ViewerContext,
+    IdentifiedViewSystem, SpaceViewSystemExecutionError, TensorDecodeCache, ViewContextCollection,
+    ViewPartSystem, ViewQuery, ViewerContext, VisualizerAdditionalApplicabilityFilter,
 };
 
 #[derive(Default)]
@@ -22,6 +21,19 @@ impl IdentifiedViewSystem for TensorSystem {
     }
 }
 
+struct TensorVisualizerEntityFilter;
+
+impl VisualizerAdditionalApplicabilityFilter for TensorVisualizerEntityFilter {
+    fn update_applicability(&mut self, event: &re_arrow_store::StoreEvent) -> bool {
+        event.diff.cells.iter().any(|(component_name, cell)| {
+            component_name == &re_types::components::TensorData::name()
+                && re_types::components::TensorData::from_arrow(cell.as_arrow_ref())
+                    .map(|tensors| tensors.iter().any(|tensor| !tensor.is_vector()))
+                    .unwrap_or(false)
+        })
+    }
+}
+
 impl ViewPartSystem for TensorSystem {
     fn required_components(&self) -> ComponentNameSet {
         Tensor::required_components()
@@ -34,26 +46,8 @@ impl ViewPartSystem for TensorSystem {
         std::iter::once(Tensor::indicator().name()).collect()
     }
 
-    fn heuristic_filter(
-        &self,
-        store: &re_arrow_store::DataStore,
-        ent_path: &EntityPath,
-        _ctx: HeuristicFilterContext,
-        query: &LatestAtQuery,
-        entity_components: &ComponentNameSet,
-    ) -> bool {
-        if !default_heuristic_filter(entity_components, &self.indicator_components()) {
-            return false;
-        }
-
-        // The tensor view can't display anything with less than two dimensions.
-        if let Some(tensor) =
-            store.query_latest_component::<re_types::components::TensorData>(ent_path, query)
-        {
-            !tensor.is_vector()
-        } else {
-            false
-        }
+    fn applicability_filter(&self) -> Option<Box<dyn VisualizerAdditionalApplicabilityFilter>> {
+        Some(Box::new(TensorVisualizerEntityFilter))
     }
 
     fn execute(

crates/re_viewer_context/src/lib.rs

@@ -49,6 +49,7 @@ pub use space_view::{
     SpaceViewOutlineMasks, SpaceViewState, SpaceViewSystemExecutionError,
     SpaceViewSystemRegistrator, SystemExecutionOutput, ViewContextCollection, ViewContextSystem,
     ViewPartCollection, ViewPartSystem, ViewQuery, ViewSystemIdentifier,
+    VisualizerAdditionalApplicabilityFilter,
 };
 pub use store_context::StoreContext;
 pub use tensor::{TensorDecodeCache, TensorStats, TensorStatsCache};

crates/re_viewer_context/src/space_view/mod.rs

@@ -33,6 +33,7 @@ pub use view_part_system::{
     default_heuristic_filter, HeuristicFilterContext, ViewPartCollection, ViewPartSystem,
 };
 pub use view_query::{DataResult, PerSystemDataResults, ViewQuery};
+pub use visualizer_entity_subscriber::VisualizerAdditionalApplicabilityFilter;
 
 // ---------------------------------------------------------------------------
 

crates/re_viewer_context/src/space_view/view_part_system.rs

@@ -7,6 +7,7 @@ use re_types::ComponentNameSet;
 use crate::{
     IdentifiedViewSystem, SpaceViewClassIdentifier, SpaceViewSystemExecutionError,
     ViewContextCollection, ViewQuery, ViewSystemIdentifier, ViewerContext,
+    VisualizerAdditionalApplicabilityFilter,
 };
 
 /// This is additional context made available to the `heuristic_filter`.
@@ -81,6 +82,13 @@ pub trait ViewPartSystem: Send + Sync + 'static {
         default_heuristic_filter(entity_components, &self.indicator_components())
     }
 
+    /// Additional filter for applicability.
+    ///
+    /// If none is specified, applicability is solely determined by required components.
+    fn applicability_filter(&self) -> Option<Box<dyn VisualizerAdditionalApplicabilityFilter>> {
+        None
+    }
+
     /// Queries the data store and performs data conversions to make it ready for display.
     ///
     /// Mustn't query any data outside of the archetype.

crates/re_viewer_context/src/space_view/visualizer_entity_subscriber.rs

@@ -28,17 +28,45 @@ pub struct VisualizerEntitySubscriber {
     required_components_indices: IntMap<ComponentName, usize>,
 
     per_store_mapping: HashMap<StoreId, VisualizerEntityMapping>,
+
+    /// Additional filter for applicability.
+    applicability_filter: Box<dyn VisualizerAdditionalApplicabilityFilter>,
+}
+
+/// Additional filter for applicability on top of the default check for required components.
+pub trait VisualizerAdditionalApplicabilityFilter: Send + Sync {
+    /// Updates the internal applicability filter state based on the given events.
+    ///
+    /// Called for every update no matter whether the entity is already has all required components or not.
+    ///
+    /// Returns true if the entity changed in the event is now applicable to the visualizer, false otherwise.
+    /// Once a entity passes this filter, it can never go back to being filtered out.
+    /// **This implies that the filter does not _need_ to be stateful.**
+    /// It is perfectly fine to return `true` only if something in the diff is regarded as applicable and false otherwise.
+    /// (However, if necessary, the applicability filter *can* keep track of state.)
+    fn update_applicability(&mut self, _event: &re_arrow_store::StoreEvent) -> bool;
+}
+
+struct DefaultVisualizerApplicabilityFilter;
+
+impl VisualizerAdditionalApplicabilityFilter for DefaultVisualizerApplicabilityFilter {
+    #[inline]
+    fn update_applicability(&mut self, _event: &re_arrow_store::StoreEvent) -> bool {
+        true
+    }
 }
 
 #[derive(Default)]
 struct VisualizerEntityMapping {
     /// For each entity, which of the required components are present.
     ///
+    /// Last bit is used for the applicability filter.
+    ///
     /// In order of `required_components`.
     /// If all bits are set, the entity is applicable to the visualizer.
     // TODO(andreas): We could just limit the number of required components to 32 or 64 and
     // then use a single u32/u64 as a bitmap.
-    required_component_bitmap_per_entity: IntMap<EntityPathHash, BitVec>,
+    required_component_and_filter_bitmap_per_entity: IntMap<EntityPathHash, BitVec>,
 
     /// Which entities the visualizer can be applied to.
     ///
@@ -59,6 +87,9 @@ impl VisualizerEntitySubscriber {
                 .map(|(i, name)| (name, i))
                 .collect(),
             per_store_mapping: Default::default(),
+            applicability_filter: visualizer
+                .applicability_filter()
+                .unwrap_or_else(|| Box::new(DefaultVisualizerApplicabilityFilter)),
         }
     }
 
@@ -102,10 +133,10 @@ impl StoreSubscriber for VisualizerEntitySubscriber {
                 .or_default();
 
             let required_components_bitmap = store_mapping
-                .required_component_bitmap_per_entity
+                .required_component_and_filter_bitmap_per_entity
                 .entry(event.diff.entity_path.hash())
                 .or_insert_with(|| {
-                    BitVec::from_elem(self.required_components_indices.len(), false)
+                    BitVec::from_elem(self.required_components_indices.len() + 1, false)
                 });
 
             if required_components_bitmap.all() {
@@ -119,6 +150,15 @@ impl StoreSubscriber for VisualizerEntitySubscriber {
                 }
             }
 
+            let bit_index_for_filter = self.required_components_indices.len();
+            let custom_filter = required_components_bitmap[bit_index_for_filter];
+            if !custom_filter {
+                required_components_bitmap.set(
+                    bit_index_for_filter,
+                    self.applicability_filter.update_applicability(event),
+                );
+            }
+
             if required_components_bitmap.all() {
                 re_log::debug!(
                     "Entity {:?} in store {:?} is now applicable to visualizer {:?}",