Chunkify time panel data density graphs (#6847)

Cargo.lock

@@ -6490,6 +6490,17 @@ dependencies = [
  "rerun",
 ]
 
+[[package]]
+name = "test_data_density_graph"
+version = "0.18.0-alpha.1+dev"
+dependencies = [
+ "anyhow",
+ "clap",
+ "rand",
+ "re_log",
+ "rerun",
+]
+
 [[package]]
 name = "test_image_memory"
 version = "0.18.0-alpha.1+dev"

crates/store/re_chunk/src/chunk.rs

@@ -221,6 +221,41 @@ impl Chunk {
             .collect()
     }
 
+    /// The cumulative number of events in this chunk.
+    ///
+    /// I.e. how many _component batches_ ("cells") were logged in total?
+    //
+    // TODO(cmc): This needs to be stored in chunk metadata and transported across IPC.
+    #[inline]
+    pub fn num_events_cumulative(&self) -> usize {
+        // Reminder: component columns are sparse, we must take a look at the validity bitmaps.
+        self.components
+            .values()
+            .map(|list_array| {
+                list_array.validity().map_or_else(
+                    || list_array.len(),
+                    |validity| validity.len() - validity.unset_bits(),
+                )
+            })
+            .sum()
+    }
+
+    /// The number of events in this chunk for the specified component.
+    ///
+    /// I.e. how many _component batches_ ("cells") were logged in total for this component?
+    //
+    // TODO(cmc): This needs to be stored in chunk metadata and transported across IPC.
+    #[inline]
+    pub fn num_events_for_component(&self, component_name: ComponentName) -> Option<usize> {
+        // Reminder: component columns are sparse, we must check validity bitmap.
+        self.components.get(&component_name).map(|list_array| {
+            list_array.validity().map_or_else(
+                || list_array.len(),
+                |validity| validity.len() - validity.unset_bits(),
+            )
+        })
+    }
+
     /// Computes the `RowId` range covered by each individual component column on each timeline.
     ///
     /// This is different from the `RowId` range covered by the [`Chunk`] as a whole because component

crates/top/re_sdk/src/recording_stream.rs

@@ -1524,6 +1524,21 @@ impl RecordingStream {
         }
     }
 
+    /// Records a single [`Chunk`].
+    ///
+    /// This will _not_ inject `log_tick` and `log_time` timeline columns into the chunk,
+    /// for that use [`Self::record_chunk`].
+    #[inline]
+    pub fn record_chunk_raw(&self, chunk: Chunk) {
+        let f = move |inner: &RecordingStreamInner| {
+            inner.batcher.push_chunk(chunk);
+        };
+
+        if self.with(f).is_none() {
+            re_log::warn_once!("Recording disabled - call to record_chunk() ignored");
+        }
+    }
+
     /// Swaps the underlying sink for a new one.
     ///
     /// This guarantees that:

crates/top/re_sdk/src/spawn.rs

@@ -45,6 +45,9 @@ pub struct SpawnOptions {
     /// Extra arguments that will be passed as-is to the Rerun Viewer process.
     pub extra_args: Vec<String>,
 
+    /// Extra environment variables that will be passed as-is to the Rerun Viewer process.
+    pub extra_env: Vec<(String, String)>,
+
     /// Hide the welcome screen.
     pub hide_welcome_screen: bool,
 }
@@ -61,6 +64,7 @@ impl Default for SpawnOptions {
             executable_name: RERUN_BINARY.into(),
             executable_path: None,
             extra_args: Vec::new(),
+            extra_env: Vec::new(),
             hide_welcome_screen: false,
         }
     }
@@ -268,6 +272,7 @@ pub fn spawn(opts: &SpawnOptions) -> Result<(), SpawnError> {
     }
 
     rerun_bin.args(opts.extra_args.clone());
+    rerun_bin.envs(opts.extra_env.clone());
 
     // SAFETY: This code is only run in the child fork, we are not modifying any memory
     // that is shared with the parent process.
@@ -290,7 +295,8 @@ pub fn spawn(opts: &SpawnOptions) -> Result<(), SpawnError> {
         // NOTE: The timeout only covers the TCP handshake: if no process is bound to that address
         // at all, the connection will fail immediately, irrelevant of the timeout configuration.
         // For that reason we use an extra loop.
-        for _ in 0..5 {
+        for i in 0..5 {
+            re_log::debug!("connection attempt {}", i + 1);
             if TcpStream::connect_timeout(&connect_addr, Duration::from_secs(1)).is_ok() {
                 break;
             }

crates/viewer/re_time_panel/src/data_density_graph.rs

@@ -3,14 +3,22 @@
 //! The data density is the number of data points per unit of time.
 //! We collect this into a histogram, blur it, and then paint it.
 
+use std::collections::HashSet;
 use std::ops::RangeInclusive;
+use std::sync::Arc;
 
 use egui::emath::Rangef;
 use egui::{epaint::Vertex, lerp, pos2, remap, Color32, NumExt as _, Rect, Shape};
 
+use re_chunk_store::Chunk;
+use re_chunk_store::RangeQuery;
 use re_data_ui::item_ui;
 use re_entity_db::TimeHistogram;
+use re_log_types::EntityPath;
+use re_log_types::TimeInt;
+use re_log_types::Timeline;
 use re_log_types::{ComponentPath, ResolvedTimeRange, TimeReal};
+use re_types::ComponentName;
 use re_viewer_context::{Item, TimeControl, UiLayout, ViewerContext};
 
 use crate::TimePanelItem;
@@ -366,6 +374,245 @@ fn smooth(density: &[f32]) -> Vec<f32> {
 
 // ----------------------------------------------------------------------------
 
+#[allow(clippy::too_many_arguments)]
+pub fn data_density_graph_ui2(
+    data_density_graph_painter: &mut DataDensityGraphPainter,
+    ctx: &ViewerContext<'_>,
+    time_ctrl: &TimeControl,
+    db: &re_entity_db::EntityDb,
+    time_area_painter: &egui::Painter,
+    ui: &egui::Ui,
+    time_ranges_ui: &TimeRangesUi,
+    row_rect: Rect,
+    item: &TimePanelItem,
+) {
+    re_tracing::profile_function!();
+
+    let timeline = *time_ctrl.timeline();
+
+    let mut data = DensityDataAggregate::new(ui, time_ranges_ui, row_rect);
+
+    // Collect all relevant chunks in the visible time range.
+    // We do this as a separate step so that we can also deduplicate chunks.
+    let visible_time_range = time_ranges_ui
+        .time_range_from_x_range((row_rect.left() - MARGIN_X)..=(row_rect.right() + MARGIN_X));
+    let mut chunk_ranges: Vec<(Arc<Chunk>, ResolvedTimeRange, usize)> = vec![];
+
+    visit_relevant_chunks(
+        db,
+        &item.entity_path,
+        item.component_name,
+        timeline,
+        visible_time_range,
+        |chunk, time_range, num_events| {
+            chunk_ranges.push((chunk, time_range, num_events));
+        },
+    );
+
+    for (_, time_range, num_events) in chunk_ranges {
+        data.add_chunk_range(time_range, num_events);
+    }
+
+    data.density_graph.buckets = smooth(&data.density_graph.buckets);
+
+    data.density_graph.paint(
+        data_density_graph_painter,
+        row_rect.y_range(),
+        time_area_painter,
+        graph_color(ctx, &item.to_item(), ui),
+        // TODO(jprochazk): completely remove `hovered_x_range` and associated code from painter
+        0f32..=0f32,
+    );
+
+    if let Some(hovered_time) = data.hovered_time {
+        ctx.selection_state().set_hovered(item.to_item());
+
+        if ui.ctx().dragged_id().is_none() {
+            // TODO(jprochazk): check chunk.num_rows() and chunk.timeline.is_sorted()
+            //                  if too many rows and unsorted, show some generic error tooltip (=too much data)
+            egui::show_tooltip_at_pointer(
+                ui.ctx(),
+                ui.layer_id(),
+                egui::Id::new("data_tooltip"),
+                |ui| {
+                    show_row_ids_tooltip2(ctx, ui, time_ctrl, db, item, hovered_time);
+                },
+            );
+        }
+    }
+}
+
+fn show_row_ids_tooltip2(
+    ctx: &ViewerContext<'_>,
+    ui: &mut egui::Ui,
+    time_ctrl: &TimeControl,
+    db: &re_entity_db::EntityDb,
+    item: &TimePanelItem,
+    at_time: TimeInt,
+) {
+    use re_data_ui::DataUi as _;
+
+    let ui_layout = UiLayout::Tooltip;
+    let query = re_chunk_store::LatestAtQuery::new(*time_ctrl.timeline(), at_time);
+
+    let TimePanelItem {
+        entity_path,
+        component_name,
+    } = item;
+
+    if let Some(component_name) = component_name {
+        let component_path = ComponentPath::new(entity_path.clone(), *component_name);
+        item_ui::component_path_button(ctx, ui, &component_path, db);
+        ui.add_space(8.0);
+        component_path.data_ui(ctx, ui, ui_layout, &query, db);
+    } else {
+        let instance_path = re_entity_db::InstancePath::entity_all(entity_path.clone());
+        item_ui::instance_path_button(ctx, &query, db, ui, None, &instance_path);
+        ui.add_space(8.0);
+        instance_path.data_ui(ctx, ui, ui_layout, &query, db);
+    }
+}
+
+struct DensityDataAggregate<'a> {
+    time_ranges_ui: &'a TimeRangesUi,
+    row_rect: Rect,
+
+    pointer_pos: Option<egui::Pos2>,
+    interact_radius: f32,
+
+    density_graph: DensityGraph,
+    hovered_time: Option<TimeInt>,
+}
+
+impl<'a> DensityDataAggregate<'a> {
+    fn new(ui: &'a egui::Ui, time_ranges_ui: &'a TimeRangesUi, row_rect: Rect) -> Self {
+        let pointer_pos = ui.input(|i| i.pointer.hover_pos());
+        let interact_radius = ui.style().interaction.resize_grab_radius_side;
+
+        Self {
+            time_ranges_ui,
+            row_rect,
+
+            pointer_pos,
+            interact_radius,
+
+            density_graph: DensityGraph::new(row_rect.x_range()),
+            hovered_time: None,
+        }
+    }
+
+    fn add_chunk_range(&mut self, time_range: ResolvedTimeRange, num_events: usize) {
+        if num_events == 0 {
+            return;
+        }
+
+        let (Some(min_x), Some(max_x)) = (
+            self.time_ranges_ui.x_from_time_f32(time_range.min().into()),
+            self.time_ranges_ui.x_from_time_f32(time_range.max().into()),
+        ) else {
+            return;
+        };
+
+        self.density_graph
+            .add_range((min_x, max_x), num_events as _);
+
+        if let Some(pointer_pos) = self.pointer_pos {
+            let is_hovered = if (max_x - min_x).abs() < 1.0 {
+                // Are we close enough to center?
+                let center_x = (max_x + min_x) / 2.0;
+                let distance_sq = pos2(center_x, self.row_rect.center().y).distance_sq(pointer_pos);
+
+                distance_sq < self.interact_radius.powi(2)
+            } else {
+                // Are we within time range rect?
+                let time_range_rect = Rect {
+                    min: egui::pos2(min_x, self.row_rect.min.y),
+                    max: egui::pos2(max_x, self.row_rect.max.y),
+                };
+
+                time_range_rect.contains(pointer_pos)
+            };
+
+            if is_hovered {
+                if let Some(at_time) = self.time_ranges_ui.time_from_x_f32(pointer_pos.x) {
+                    self.hovered_time = Some(at_time.round());
+                }
+            }
+        }
+    }
+}
+
+/// This is a wrapper over `range_relevant_chunks` which also supports querying the entire entity.
+/// Relevant chunks are those which:
+/// - Contain data for `entity_path`
+/// - Contain a `component_name` column (if provided)
+/// - Have data on the given `timeline`
+/// - Have data in the given `time_range`
+///
+/// The does not deduplicates chunks when no `component_name` is provided.
+fn visit_relevant_chunks(
+    db: &re_entity_db::EntityDb,
+    entity_path: &EntityPath,
+    component_name: Option<ComponentName>,
+    timeline: Timeline,
+    time_range: ResolvedTimeRange,
+    mut visitor: impl FnMut(Arc<Chunk>, ResolvedTimeRange, usize),
+) {
+    re_tracing::profile_function!();
+
+    let query = RangeQuery::new(timeline, time_range);
+
+    if let Some(component_name) = component_name {
+        let chunks = db
+            .store()
+            .range_relevant_chunks(&query, entity_path, component_name);
+
+        for chunk in chunks {
+            let Some(num_events) = chunk.num_events_for_component(component_name) else {
+                continue;
+            };
+
+            let Some(chunk_timeline) = chunk.timelines().get(&timeline) else {
+                continue;
+            };
+
+            visitor(Arc::clone(&chunk), chunk_timeline.time_range(), num_events);
+        }
+    } else {
+        let mut seen = HashSet::new();
+        if let Some(subtree) = db.tree().subtree(entity_path) {
+            subtree.visit_children_recursively(&mut |entity_path, _| {
+                let Some(components) = db.store().all_components(&timeline, entity_path) else {
+                    return;
+                };
+
+                for component_name in components {
+                    let chunks =
+                        db.store()
+                            .range_relevant_chunks(&query, entity_path, component_name);
+
+                    for chunk in chunks {
+                        let Some(chunk_timeline) = chunk.timelines().get(&timeline) else {
+                            continue;
+                        };
+
+                        if seen.contains(&chunk.id()) {
+                            continue;
+                        }
+                        seen.insert(chunk.id());
+
+                        visitor(
+                            Arc::clone(&chunk),
+                            chunk_timeline.time_range(),
+                            chunk.num_events_cumulative(),
+                        );
+                    }
+                }
+            });
+        }
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn data_density_graph_ui(
     data_density_graph_painter: &mut DataDensityGraphPainter,
@@ -545,7 +792,7 @@ fn show_row_ids_tooltip(
     }
 
     let ui_layout = UiLayout::Tooltip;
-    let query = re_chunk_store::LatestAtQuery::new(*time_ctrl.timeline(), time_range.max());
+    let query = re_chunk_store::LatestAtQuery::new(*time_ctrl.timeline(), time_range.center());
 
     let TimePanelItem {
         entity_path,