Time panel chunkification

crates/store/re_chunk_store/src/query.rs

@@ -5,7 +5,9 @@ use std::{
 
 use itertools::Itertools;
 use re_chunk::{Chunk, LatestAtQuery, RangeQuery};
+use re_log_types::ResolvedTimeRange;
 use re_log_types::{EntityPath, TimeInt, Timeline};
+use re_types_core::SizeBytes as _;
 use re_types_core::{ComponentName, ComponentNameSet};
 
 use crate::{store::ChunkIdSetPerTime, ChunkStore};
@@ -61,6 +63,47 @@ impl ChunkStore {
         }
     }
 
+    /// Retrieve all the [`ComponentName`]s that have been written to for a given [`EntityPath`].
+    ///
+    /// Static components are always included in the results.
+    ///
+    /// Returns `None` if the entity has never had any data logged to it.
+    pub fn all_components_on_all_timelines(
+        &self,
+        entity_path: &EntityPath,
+    ) -> Option<ComponentNameSet> {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        let static_components: Option<ComponentNameSet> = self
+            .static_chunk_ids_per_entity
+            .get(entity_path)
+            .map(|static_chunks_per_component| {
+                static_chunks_per_component.keys().copied().collect()
+            });
+
+        let temporal_components: Option<ComponentNameSet> = self
+            .temporal_chunk_ids_per_entity_per_component
+            .get(entity_path)
+            .map(|temporal_chunk_ids_per_timeline| {
+                temporal_chunk_ids_per_timeline
+                    .iter()
+                    .flat_map(|(_, temporal_chunk_ids_per_component)| {
+                        temporal_chunk_ids_per_component.keys().copied()
+                    })
+                    .collect()
+            });
+
+        match (static_components, temporal_components) {
+            (None, None) => None,
+            (None, comps @ Some(_)) | (comps @ Some(_), None) => comps,
+            (Some(static_comps), Some(temporal_comps)) => {
+                Some(static_comps.into_iter().chain(temporal_comps).collect())
+            }
+        }
+    }
+
     /// Check whether a given entity has a specific [`ComponentName`] either on the specified
     /// timeline, or in its static data.
     #[inline]
@@ -75,6 +118,56 @@ impl ChunkStore {
             .map_or(false, |components| components.contains(component_name))
     }
 
+    pub fn entity_has_component_on_any_timeline(
+        &self,
+        entity_path: &EntityPath,
+        component_name: &ComponentName,
+    ) -> bool {
+        re_tracing::profile_function!();
+
+        if self
+            .static_chunk_ids_per_entity
+            .get(entity_path)
+            .and_then(|static_chunks_per_component| static_chunks_per_component.get(component_name))
+            .and_then(|id| self.chunks_per_chunk_id.get(id))
+            .is_some()
+        {
+            return true;
+        }
+
+        for temporal_chunk_ids_per_component in self
+            .temporal_chunk_ids_per_entity_per_component
+            .get(entity_path)
+            .iter()
+            .flat_map(|temporal_chunk_ids_per_timeline| temporal_chunk_ids_per_timeline.values())
+        {
+            if temporal_chunk_ids_per_component
+                .get(component_name)
+                .is_some()
+            {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    /// Returns true if the given entity has a specific component with static data.
+    #[inline]
+    pub fn entity_has_static_component(
+        &self,
+        entity_path: &EntityPath,
+        component_name: &ComponentName,
+    ) -> bool {
+        re_tracing::profile_function!();
+
+        let Some(static_components) = self.static_chunk_ids_per_entity.get(entity_path) else {
+            return false;
+        };
+
+        static_components.contains_key(component_name)
+    }
+
     /// Find the earliest time at which something was logged for a given entity on the specified
     /// timeline.
     ///
@@ -413,3 +506,217 @@ impl ChunkStore {
             .collect()
     }
 }
+
+// Queries returning `usize` and `bool`
+impl ChunkStore {
+    /// Returns the number of events logged for a given component on the given entity path across all timelines.
+    pub fn num_events_on_timeline_for_component(
+        &self,
+        timeline: &Timeline,
+        entity_path: &EntityPath,
+        component_name: ComponentName,
+    ) -> usize {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        self.num_events_for_static_component(entity_path, component_name)
+            + self.num_events_on_timeline_for_temporal_component(
+                timeline,
+                entity_path,
+                component_name,
+            )
+    }
+
+    pub fn time_range_for_entity(
+        &self,
+        timeline: &Timeline,
+        entity_path: &EntityPath,
+    ) -> Option<ResolvedTimeRange> {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        let temporal_chunk_ids_per_timeline =
+            self.temporal_chunk_ids_per_entity.get(entity_path)?;
+        let chunk_id_sets = temporal_chunk_ids_per_timeline.get(timeline)?;
+
+        let start = chunk_id_sets.per_start_time.keys().min()?;
+        let end = chunk_id_sets.per_end_time.keys().max()?;
+
+        Some(ResolvedTimeRange::new(*start, *end))
+    }
+
+    pub fn num_events_on_timeline_for_all_components(
+        &self,
+        timeline: &Timeline,
+        entity_path: &EntityPath,
+    ) -> usize {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        let mut total_events = 0;
+
+        if let Some(static_chunks_per_component) = self.static_chunk_ids_per_entity.get(entity_path)
+        {
+            for chunk in static_chunks_per_component
+                .values()
+                .filter_map(|id| self.chunks_per_chunk_id.get(id))
+            {
+                total_events += chunk.num_events_cumulative();
+            }
+        }
+
+        if let Some(chunk_ids) = self
+            .temporal_chunk_ids_per_entity
+            .get(entity_path)
+            .and_then(|temporal_chunks_events_per_timeline| {
+                temporal_chunks_events_per_timeline.get(timeline)
+            })
+        {
+            for chunk in chunk_ids
+                .per_start_time
+                .values()
+                .flat_map(|ids| ids.iter().filter_map(|id| self.chunks_per_chunk_id.get(id)))
+            {
+                total_events += chunk.num_events_cumulative();
+            }
+        }
+
+        total_events
+    }
+
+    pub fn entity_has_data_on_timeline(
+        &self,
+        timeline: &Timeline,
+        entity_path: &EntityPath,
+    ) -> bool {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        if self.static_chunk_ids_per_entity.get(entity_path).is_some() {
+            // Static data exists on all timelines
+            return true;
+        }
+
+        if let Some(temporal_chunk_ids_per_timeline) =
+            self.temporal_chunk_ids_per_entity.get(entity_path)
+        {
+            if temporal_chunk_ids_per_timeline.contains_key(timeline) {
+                return true;
+            }
+        }
+
+        false
+    }
+
+    pub fn num_events_for_static_component(
+        &self,
+        entity_path: &EntityPath,
+        component_name: ComponentName,
+    ) -> usize {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        if let Some(static_chunk) = self
+            .static_chunk_ids_per_entity
+            .get(entity_path)
+            .and_then(|static_chunks_per_component| {
+                static_chunks_per_component.get(&component_name)
+            })
+            .and_then(|chunk_id| self.chunks_per_chunk_id.get(chunk_id))
+        {
+            // Static data overrides temporal data
+            static_chunk
+                .num_events_for_component(component_name)
+                .unwrap_or(0)
+        } else {
+            0
+        }
+    }
+
+    /// Returns the number of events logged for a given component on the given entity path across all timelines.
+    ///
+    /// This ignores static data.
+    pub fn num_events_on_timeline_for_temporal_component(
+        &self,
+        timeline: &Timeline,
+        entity_path: &EntityPath,
+        component_name: ComponentName,
+    ) -> usize {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        let Some(temporal_chunk_ids_per_timeline) = self
+            .temporal_chunk_ids_per_entity_per_component
+            .get(entity_path)
+        else {
+            return 0; // no events logged for the entity path
+        };
+
+        let Some(temporal_chunk_ids_per_component) = temporal_chunk_ids_per_timeline.get(timeline)
+        else {
+            return 0; // no events logged on this timeline
+        };
+
+        let Some(chunk_ids) = temporal_chunk_ids_per_component.get(&component_name) else {
+            return 0; // no events logged for the component on this timeline
+        };
+
+        let mut num_events = 0;
+        for chunk in chunk_ids
+            .per_start_time
+            .values()
+            .flat_map(|ids| ids.iter().filter_map(|id| self.chunks_per_chunk_id.get(id)))
+        {
+            num_events += chunk.num_events_for_component(component_name).unwrap_or(0);
+        }
+
+        num_events
+    }
+
+    pub fn size_of_entity_on_timeline(
+        &self,
+        timeline: &Timeline,
+        entity_path: &EntityPath,
+    ) -> usize {
+        re_tracing::profile_function!();
+
+        self.query_id.fetch_add(1, Ordering::Relaxed);
+
+        let mut total_size = 0;
+
+        if let Some(static_chunks_per_component) = self.static_chunk_ids_per_entity.get(entity_path)
+        {
+            for chunk in static_chunks_per_component
+                .values()
+                .filter_map(|id| self.chunks_per_chunk_id.get(id))
+            {
+                total_size += Chunk::total_size_bytes(chunk) as usize;
+            }
+        }
+
+        if let Some(chunk_id_sets) = self
+            .temporal_chunk_ids_per_entity
+            .get(entity_path)
+            .and_then(|temporal_chunk_ids_per_timeline| {
+                temporal_chunk_ids_per_timeline.get(timeline)
+            })
+        {
+            for chunk in chunk_id_sets
+                .per_start_time
+                .values()
+                .flat_map(|v| v.iter())
+                .filter_map(|id| self.chunks_per_chunk_id.get(id))
+            {
+                total_size += Chunk::total_size_bytes(chunk) as usize;
+            }
+        }
+
+        total_size
+    }
+}