Add support to expand rows to show instances in the dataframe view

Cargo.lock

@@ -1947,7 +1947,7 @@ dependencies = [
 [[package]]
 name = "egui_table"
 version = "0.28.1"
-source = "git+https://github.com/rerun-io/egui_table.git?rev=0f594701d528c4a9553521cb941de1886549dc70#0f594701d528c4a9553521cb941de1886549dc70"
+source = "git+https://github.com/rerun-io/egui_table.git?rev=1c5db2cc3c4c8db79e85824f7ca6cca87110111f#1c5db2cc3c4c8db79e85824f7ca6cca87110111f"
 dependencies = [
  "egui",
  "serde",

crates/viewer/re_space_view_dataframe/Cargo.toml

@@ -36,7 +36,7 @@ re_viewport_blueprint.workspace = true
 
 anyhow.workspace = true
 egui_extras.workspace = true
-egui_table = { git = "https://github.com/rerun-io/egui_table.git", rev = "0f594701d528c4a9553521cb941de1886549dc70" } # main as of 2024-09-09
+egui_table = { git = "https://github.com/rerun-io/egui_table.git", rev = "1c5db2cc3c4c8db79e85824f7ca6cca87110111f" } # main as of 2024-09-12
 egui.workspace = true
 itertools.workspace = true
 thiserror.workspace = true

crates/viewer/re_space_view_dataframe/src/dataframe_ui.rs

@@ -13,14 +13,16 @@
 use re_viewer_context::ViewerContext;
 
 use crate::display_record_batch::{DisplayRecordBatch, DisplayRecordBatchError};
+use crate::expanded_rows::{ExpandedRows, ExpandedRowsCache, QueryExpression};
 
 /// Display a dataframe table for the provided query.
 pub(crate) fn dataframe_ui<'a>(
     ctx: &ViewerContext<'_>,
     ui: &mut egui::Ui,
     query: impl Into<QueryHandle<'a>>,
+    expanded_rows_cache: &mut ExpandedRowsCache,
 ) {
-    dataframe_ui_impl(ctx, ui, &query.into());
+    dataframe_ui_impl(ctx, ui, &query.into(), expanded_rows_cache);
 }
 
 /// A query handle for either a latest-at or range query.
@@ -62,6 +64,13 @@
             QueryHandle::Range(query_handle) => query_handle.query().timeline,
         }
     }
+
+    fn query_expression(&self) -> QueryExpression {
+        match self {
+            QueryHandle::LatestAt(query_handle) => query_handle.query().clone().into(),
+            QueryHandle::Range(query_handle) => query_handle.query().clone().into(),
+        }
+    }
 }
 
 impl<'a> From<LatestAtQueryHandle<'a>> for QueryHandle<'a> {
@@ -165,6 +174,8 @@
     header_entity_paths: Vec<Option<EntityPath>>,
     display_data: anyhow::Result<RowsDisplayData>,
 
+    expanded_rows: ExpandedRows<'a>,
+
     num_rows: u64,
 }
 
@@ -239,6 +250,11 @@
 
         debug_assert!(cell.row_nr < self.num_rows, "Bug in egui_table");
 
+        // best effort to get a proper column autosizing behaviour
+        if ui.is_sizing_pass() {
+            ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Extend);
+        }
+
         let display_data = match &self.display_data {
             Ok(display_data) => display_data,
             Err(err) => {
@@ -247,6 +263,7 @@
             }
         };
 
+        // TODO(ab): this is getting wild and should be refactored
         let cell_ui = |ui: &mut egui::Ui| {
             if let Some(BatchRef { batch_idx, row_idx }) =
                 display_data.batch_ref_from_row.get(&cell.row_nr).copied()
@@ -276,7 +293,93 @@
                 } else {
                     ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Truncate);
                 }
-                column.data_ui(self.ctx, ui, row_id, &latest_at_query, row_idx);
+
+                let instance_count = column.instance_count(row_idx);
+                let row_expansion = self.expanded_rows.row_expansion(cell.row_nr);
+
+                let instance_indices = std::iter::once(None)
+                    .chain((0..instance_count).map(Option::Some))
+                    .take(row_expansion as usize + 1);
+
+                let sub_cell_content =
+                    |ui: &mut egui::Ui,
+                     expanded_rows: &mut ExpandedRows<'_>,
+                     sub_cell_index: usize,
+                     instance_index: Option<u64>| {
+                        if instance_index.is_none() && instance_count > 1 {
+                            let cell_clicked =
+                                cell_with_hover_button_ui(ui, Some(&re_ui::icons::EXPAND), |ui| {
+                                    ui.label(format!("{instance_count} instances"));
+                                });
+
+                            if cell_clicked {
+                                if instance_count == row_expansion {
+                                    expanded_rows.collapse_row(cell.row_nr);
+                                } else {
+                                    expanded_rows.set_row_expansion(cell.row_nr, instance_count);
+                                }
+                            }
+                        } else {
+                            let has_collapse_button = instance_index
+                                .is_some_and(|instance_index| instance_index < instance_count);
+
+                            let remaining_instances = if sub_cell_index as u64 == row_expansion {
+                                instance_index.and_then(|instance_index| {
+                                    let remaining = instance_count
+                                        .saturating_sub(instance_index)
+                                        .saturating_sub(1);
+                                    if remaining > 0 {
+                                        // +1 is because the "X more…" line takes one instance spot
+                                        Some(remaining + 1)
+                                    } else {
+                                        None
+                                    }
+                                })
+                            } else {
+                                None
+                            };
+
+                            if let Some(remaining_instances) = remaining_instances {
+                                let cell_clicked = cell_with_hover_button_ui(
+                                    ui,
+                                    Some(&re_ui::icons::EXPAND),
+                                    |ui| {
+                                        ui.label(format!("{remaining_instances} more…"));
+                                    },
+                                );
+
+                                if cell_clicked {
+                                    expanded_rows.set_row_expansion(cell.row_nr, instance_count);
+                                }
+                            } else {
+                                let cell_clicked = cell_with_hover_button_ui(
+                                    ui,
+                                    has_collapse_button.then_some(&re_ui::icons::COLLAPSE),
+                                    |ui| {
+                                        column.data_ui(
+                                            self.ctx,
+                                            ui,
+                                            row_id,
+                                            &latest_at_query,
+                                            row_idx,
+                                            instance_index,
+                                        );
+                                    },
+                                );
+
+                                if cell_clicked {
+                                    expanded_rows.collapse_row(cell.row_nr);
+                                }
+                            }
+                        }
+                    };
+
+                sub_cell_ui(
+                    ui,
+                    &mut self.expanded_rows,
+                    instance_indices,
+                    sub_cell_content,
+                );
             } else {
                 error_ui(
                     ui,
@@ -289,12 +392,29 @@
             .inner_margin(egui::Margin::symmetric(Self::LEFT_RIGHT_MARGIN, 0.0))
             .show(ui, cell_ui);
     }
+
+    fn row_top_offset(&self, _ctx: &egui::Context, _table_id: egui::Id, row_nr: u64) -> f32 {
+        self.expanded_rows.row_top_offset(row_nr)
+    }
+
+    fn default_row_height(&self) -> f32 {
+        re_ui::DesignTokens::table_line_height()
+    }
 }
 
 /// Display the result of a [`QueryHandle`] in a table.
-fn dataframe_ui_impl(ctx: &ViewerContext<'_>, ui: &mut egui::Ui, query_handle: &QueryHandle<'_>) {
+fn dataframe_ui_impl(
+    ctx: &ViewerContext<'_>,
+    ui: &mut egui::Ui,
+    query_handle: &QueryHandle<'_>,
+    expanded_rows_cache: &mut ExpandedRowsCache,
+) {
     re_tracing::profile_function!();
 
+    // Salt everything against the query expression, as different queries may lead to entirely
+    // different tables, etc.
+    let id_salt = egui::Id::new("__dataframe__").with(query_handle.query_expression());
+
     let schema = query_handle.schema();
     let (header_groups, header_entity_paths) = column_groups_for_entity(schema);
 
@@ -309,6 +429,13 @@
         display_data: Err(anyhow::anyhow!(
             "No row data, `fetch_columns_and_rows` not called."
         )),
+        expanded_rows: ExpandedRows::new(
+            ui.ctx().clone(),
+            ui.make_persistent_id(id_salt).with("expanded_rows"),
+            expanded_rows_cache,
+            query_handle.query_expression(),
+            re_ui::DesignTokens::table_line_height(),
+        ),
     };
 
     let num_sticky_cols = schema
@@ -318,6 +445,7 @@
 
     egui::Frame::none().inner_margin(5.0).show(ui, |ui| {
         egui_table::Table::new()
+            .id_salt(id_salt)
             .columns(
                 schema
                     .iter()
@@ -337,7 +465,6 @@
                 egui_table::HeaderRow::new(re_ui::DesignTokens::table_header_height()),
             ])
             .num_rows(num_rows)
-            .row_height(re_ui::DesignTokens::table_line_height())
             .show(ui, &mut table_delegate);
     });
 }
@@ -375,3 +502,108 @@
     ui.error_label(error);
     re_log::warn_once!("{error}");
 }
+
+/// Draw some cell content with an optional, right-aligned, on-hover button.
+///
+/// If no icon is provided, no button is shown. Returns true if the button was shown and the cell
+/// was clicked.
+// TODO(ab, emilk): ideally, egui::Sides should work for that, but it doesn't yet support the
+// symmetric behavior (left variable width, right fixed width).
+fn cell_with_hover_button_ui(
+    ui: &mut egui::Ui,
+    icon: Option<&'static re_ui::Icon>,
+    cell_content: impl FnOnce(&mut egui::Ui),
+) -> bool {
+    let Some(icon) = icon else {
+        cell_content(ui);
+        return false;
+    };
+
+    if ui.is_sizing_pass() {
+        // we don't need space for the icon since it only shows on hover
+        cell_content(ui);
+        return false;
+    }
+
+    let (is_hovering_cell, is_clicked) = ui.input(|i| {
+        (
+            i.pointer
+                .interact_pos()
+                .is_some_and(|pos| ui.max_rect().contains(pos)),
+            i.pointer.button_clicked(egui::PointerButton::Primary),
+        )
+    });
+
+    if is_hovering_cell {
+        let mut content_rect = ui.max_rect();
+        content_rect.max.x = (content_rect.max.x
+            - re_ui::DesignTokens::small_icon_size().x
+            - re_ui::DesignTokens::text_to_icon_padding())
+        .at_least(content_rect.min.x);
+
+        let button_rect = egui::Rect::from_x_y_ranges(
+            (content_rect.max.x + re_ui::DesignTokens::text_to_icon_padding())
+                ..=ui.max_rect().max.x,
+            ui.max_rect().y_range(),
+        );
+
+        let mut content_ui = ui.new_child(egui::UiBuilder::new().max_rect(content_rect));
+        cell_content(&mut content_ui);
+
+        let mut button_ui = ui.new_child(egui::UiBuilder::new().max_rect(button_rect));
+        button_ui.visuals_mut().widgets.hovered.weak_bg_fill = egui::Color32::TRANSPARENT;
+        button_ui.visuals_mut().widgets.active.weak_bg_fill = egui::Color32::TRANSPARENT;
+        button_ui.add(egui::Button::image(
+            icon.as_image()
+                .fit_to_exact_size(re_ui::DesignTokens::small_icon_size())
+                .tint(button_ui.visuals().widgets.noninteractive.text_color()),
+        ));
+
+        is_clicked
+    } else {
+        cell_content(ui);
+        false
+    }
+}
+
+/// Helper to draw individual rows (aka sub-cell) into an expanded cell in a table.
+///
+/// `context`: whatever mutable context is necessary for the `sub_cell_content`
+/// `sub_cell_data`: the data to be displayed in each sub-cell
+/// `sub_cell_content`: the function to draw the content of each sub-cell
+fn sub_cell_ui<I, Ctx>(
+    ui: &mut egui::Ui,
+    context: &mut Ctx,
+    sub_cell_data: impl Iterator<Item = I>,
+    sub_cell_content: impl Fn(&mut egui::Ui, &mut Ctx, usize, I),
+) {
+    for (sub_cell_index, item_data) in sub_cell_data.enumerate() {
+        let sub_cell_rect = egui::Rect::from_min_size(
+            ui.cursor().min
+                + egui::vec2(
+                    0.0,
+                    sub_cell_index as f32 * re_ui::DesignTokens::table_line_height(),
+                ),
+            egui::vec2(
+                ui.available_width(),
+                re_ui::DesignTokens::table_line_height(),
+            ),
+        );
+
+        // During animation, there may be more sub-cells than can possibly fit. If so, no point in
+        // continuing to draw them.
+        if !ui.max_rect().intersects(sub_cell_rect) {
+            return;
+        }
+
+        if !ui.is_rect_visible(sub_cell_rect) {
+            continue;
+            // TODO(ab): detect that we are below any possible screen, in which case we can stop
+            // entirely.
+        }
+
+        let mut sub_cell_ui = ui.new_child(egui::UiBuilder::new().max_rect(sub_cell_rect));
+
+        sub_cell_content(&mut sub_cell_ui, context, sub_cell_index, item_data);
+    }
+}