update to AnyColumnDescriptor

rerun-io · abey79 · Feb 17, 2025 · Feb 10, 2025 · Feb 10, 2025 · Feb 10, 2025
commit c8538df17adfe8b6591e36c5f4b698d22edd65eb
diff --git a/Cargo.lock b/Cargo.lock
@@ -6912,6 +6912,7 @@ dependencies = [
  "re_log",
  "re_log_types",
  "re_renderer",
+ "re_sorbet",
  "re_tracing",
  "re_types",
  "re_types_core",

diff --git a/crates/viewer/re_catalog_hub/src/collection_ui.rs b/crates/viewer/re_catalog_hub/src/collection_ui.rs
@@ -1,7 +1,6 @@
-use egui::Ui;
 use egui_table::{CellInfo, HeaderCellInfo};
 
-use re_sorbet::ColumnDescriptor;
+use re_sorbet::AnyColumnDescriptor;
 use re_ui::UiExt;
 use re_view_dataframe::display_record_batch::DisplayRecordBatch;
 use re_viewer_context::external::re_log_types::Timeline;
@@ -34,24 +33,18 @@ pub fn collection_ui(
         .map(|record_batch| record_batch.num_rows() as u64)
         .sum();
 
-    let columns = sorbet_schema
-        .columns
-        .indices
-        .iter()
-        .map(|desc| ColumnDescriptor::Time(desc.clone()))
-        .chain(
-            sorbet_schema
-                .columns
-                .components
-                .iter()
-                .map(|desc| ColumnDescriptor::Component(desc.clone())),
-        )
-        .collect::<Vec<_>>();
+    let columns = sorbet_schema.columns.descriptors().collect::<Vec<_>>();
 
     let display_record_batches: Result<Vec<_>, _> = collection
         .collection
         .iter()
-        .map(|sorbet_batch| DisplayRecordBatch::try_new(&sorbet_batch.columns()[1..], &columns))
+        .map(|sorbet_batch| {
+            DisplayRecordBatch::try_new(
+                sorbet_batch
+                    .all_columns()
+                    .map(|(desc, array)| (desc, array.clone())),
+            )
+        })
         .collect();
 
     let display_record_batches = match display_record_batches {
@@ -99,21 +92,28 @@ pub fn collection_ui(
 struct CollectionTableDelegate<'a> {
     ctx: &'a ViewerContext<'a>,
     display_record_batches: &'a Vec<DisplayRecordBatch>,
-    selected_columns: &'a Vec<ColumnDescriptor>,
+    selected_columns: &'a Vec<AnyColumnDescriptor>,
 }
 
 impl egui_table::TableDelegate for CollectionTableDelegate<'_> {
-    fn header_cell_ui(&mut self, ui: &mut Ui, cell: &HeaderCellInfo) {
+    fn header_cell_ui(&mut self, ui: &mut egui::Ui, cell: &HeaderCellInfo) {
+        ui.set_truncate_style();
+
         let name = self.selected_columns[cell.group_index].short_name();
-        // let name = name.strip_prefix("rerun_").unwrap_or(name);
+        let name = name
+            .strip_prefix("rerun_")
+            .unwrap_or(name.as_str())
+            .replace('_', " ");
 
         ui.strong(name);
     }
 
-    fn cell_ui(&mut self, ui: &mut Ui, cell: &CellInfo) {
+    fn cell_ui(&mut self, ui: &mut egui::Ui, cell: &CellInfo) {
         // find record batch
         let mut row_index = cell.row_nr as usize;
 
+        ui.set_truncate_style();
+
         for display_record_batch in self.display_record_batches {
             let row_count = display_record_batch.num_rows();
             if row_index < row_count {
@@ -130,12 +130,6 @@ impl egui_table::TableDelegate for CollectionTableDelegate<'_> {
                     row_index,
                     None,
                 );
-
-                // if column.is_null(row_index) {
-                //     ui.label("-");
-                // } else {
-                //     re_ui::arrow_ui(ui, UiLayout::List, &column.slice(row_index, 1));
-                // }
             } else {
                 row_index -= row_count;
             }

diff --git a/crates/viewer/re_catalog_hub/src/hub.rs b/crates/viewer/re_catalog_hub/src/hub.rs
@@ -31,6 +31,8 @@ impl Catalog {
 pub struct RecordingCollection {
     pub collection_id: egui::Id,
 
+    pub name: String,
+
     //TODO(ab): in the future, this will be `SorbetBatch` instead
     pub collection: Vec<SorbetBatch>,
 }
@@ -142,7 +144,7 @@ impl CatalogHub {
         if catalog.is_empty() {
             ui.list_item_flat_noninteractive(list_item::LabelContent::new("(empty)").italics(true));
         } else {
-            for (index, _collection) in catalog.collections.iter().enumerate() {
+            for (index, collection) in catalog.collections.iter().enumerate() {
                 let is_selected =
                     if let Some(selected_collection) = self.selected_collection.as_ref() {
                         selected_collection.server_origin == *origin
@@ -151,7 +153,7 @@ impl CatalogHub {
                         false
                     };
 
-                let content = list_item::LabelContent::new(format!("collection {index}"));
+                let content = list_item::LabelContent::new(collection.name.as_ref());
                 let response = ui.list_item().selected(is_selected).show_flat(ui, content);
 
                 if response.clicked() {
@@ -164,6 +166,14 @@ impl CatalogHub {
                 }
             }
         }
+
+        // deselect when clicking in empty space
+        let empty_space_response = ui.allocate_response(ui.available_size(), egui::Sense::click());
+
+        // clear selection upon clicking the empty space
+        if empty_space_response.clicked() {
+            self.command_queue.lock().push(Command::DeselectCollection);
+        }
     }
 
     pub fn ui(&self, ctx: &ViewerContext<'_>, ui: &mut egui::Ui) {
@@ -240,6 +250,8 @@ async fn stream_catalog_async(
     let catalog = Catalog {
         collections: vec![RecordingCollection {
             collection_id,
+            //TODO(ab): this should be provided by the server
+            name: "default".to_owned(),
             collection: sorbet_batches,
         }],
     };

diff --git a/crates/viewer/re_ui/src/ui_ext.rs b/crates/viewer/re_ui/src/ui_ext.rs
@@ -1168,6 +1168,17 @@ pub trait UiExt {
         })
         .inner
     }
+
+    /// Set [`egui::Style::wrap_mode`] to [`egui::TextWrapMode::Truncate`], unless this is a sizing
+    /// pass, in which case it is set to [`egui::TextWrapMode::Extend`].
+    fn set_truncate_style(&mut self) {
+        let ui = self.ui_mut();
+        if ui.is_sizing_pass() {
+            ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Extend);
+        } else {
+            ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Truncate);
+        }
+    }
 }
 
 impl UiExt for egui::Ui {

diff --git a/crates/viewer/re_view_dataframe/Cargo.toml b/crates/viewer/re_view_dataframe/Cargo.toml
@@ -31,6 +31,7 @@ re_format.workspace = true
 re_log_types.workspace = true
 re_log.workspace = true
 re_renderer.workspace = true
+re_sorbet.workspace = true
 re_tracing.workspace = true
 re_types_core.workspace = true
 re_types.workspace = true

diff --git a/crates/viewer/re_view_dataframe/src/dataframe_ui.rs b/crates/viewer/re_view_dataframe/src/dataframe_ui.rs
@@ -148,7 +148,14 @@ impl RowsDisplayData {
     ) -> Result<Self, DisplayRecordBatchError> {
         let display_record_batches = row_data
             .into_iter()
-            .map(|data| DisplayRecordBatch::try_new(&data, selected_columns))
+            .map(|data| {
+                DisplayRecordBatch::try_new(
+                    selected_columns
+                        .iter()
+                        .map(|desc| desc.clone().into())
+                        .zip(data),
+                )
+            })
             .collect::<Result<Vec<_>, _>>()?;
 
         let mut batch_ref_from_row = BTreeMap::new();
@@ -226,11 +233,7 @@ impl egui_table::TableDelegate for DataframeTableDelegate<'_> {
     }
 
     fn header_cell_ui(&mut self, ui: &mut egui::Ui, cell: &egui_table::HeaderCellInfo) {
-        if ui.is_sizing_pass() {
-            ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Extend);
-        } else {
-            ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Truncate);
-        }
+        ui.set_truncate_style();
 
         egui::Frame::new()
             .inner_margin(egui::Margin::symmetric(4, 0))
@@ -410,11 +413,7 @@ impl egui_table::TableDelegate for DataframeTableDelegate<'_> {
             .unwrap_or_else(|| Timeline::new("", TimeType::Sequence));
         let latest_at_query = LatestAtQuery::new(filtered_index, timestamp);
 
-        if ui.is_sizing_pass() {
-            ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Extend);
-        } else {
-            ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Truncate);
-        }
+        ui.set_truncate_style();
 
         let instance_count = column.instance_count(batch_row_idx);
         let additional_lines = self.expanded_rows.additional_lines_for_row(cell.row_nr);

diff --git a/crates/viewer/re_view_dataframe/src/display_record_batch.rs b/crates/viewer/re_view_dataframe/src/display_record_batch.rs
@@ -1,22 +1,24 @@
 //! Intermediate data structures to make `re_datastore`'s row data more amenable to displaying in a
 //! table.
 
-use arrow::buffer::ScalarBuffer as ArrowScalarBuffer;
 use arrow::{
     array::{
         Array as ArrowArray, ArrayRef as ArrowArrayRef,
         Int32DictionaryArray as ArrowInt32DictionaryArray, ListArray as ArrowListArray,
     },
     buffer::NullBuffer as ArrowNullBuffer,
+    buffer::ScalarBuffer as ArrowScalarBuffer,
     datatypes::DataType as ArrowDataType,
 };
 use thiserror::Error;
 
 use re_arrow_util::ArrowArrayDowncastRef as _;
-use re_chunk_store::{ColumnDescriptor, ComponentColumnDescriptor, LatestAtQuery};
+use re_chunk_store::LatestAtQuery;
 use re_dataframe::external::re_chunk::{TimeColumn, TimeColumnError};
+use re_log_types::external::re_tuid::Tuid;
 use re_log_types::{EntityPath, TimeInt, Timeline};
-use re_types_core::ComponentName;
+use re_sorbet::{AnyColumnDescriptor, ComponentColumnDescriptor};
+use re_types_core::{ComponentName, DeserializationError, Loggable};
 use re_ui::UiExt;
 use re_viewer_context::{UiLayout, ViewerContext};
 
@@ -30,11 +32,14 @@ pub enum DisplayRecordBatchError {
 
     #[error("Unexpected column data type for component '{0}': {1:?}")]
     UnexpectedComponentColumnDataType(String, ArrowDataType),
+
+    #[error(transparent)]
+    DeserializationError(#[from] DeserializationError),
 }
 
 /// A single column of component data.
 ///
-/// Abstracts over the different possible arrow representation of component data.
+/// Abstracts over the different possible arrow representations of component data.
 #[derive(Debug)]
 pub enum ComponentData {
     Null,
@@ -166,6 +171,9 @@ impl ComponentData {
 /// A single column of data in a record batch.
 #[derive(Debug)]
 pub enum DisplayColumn {
+    RowId {
+        row_ids: Vec<Tuid>,
+    },
     Timeline {
         timeline: Timeline,
         time_data: ArrowScalarBuffer<i64>,
@@ -180,11 +188,15 @@ pub enum DisplayColumn {
 
 impl DisplayColumn {
     fn try_new(
-        column_descriptor: &ColumnDescriptor,
+        column_descriptor: &AnyColumnDescriptor,
         column_data: &ArrowArrayRef,
     ) -> Result<Self, DisplayRecordBatchError> {
         match column_descriptor {
-            ColumnDescriptor::Time(desc) => {
+            AnyColumnDescriptor::RowId(_desc) => Ok(Self::RowId {
+                row_ids: Tuid::from_arrow(column_data)?,
+            }),
+
+            AnyColumnDescriptor::Time(desc) => {
                 let timeline = desc.timeline();
 
                 let (time_data, time_nulls) = TimeColumn::read_nullable_array(column_data)
@@ -199,8 +211,7 @@ impl DisplayColumn {
                     time_nulls,
                 })
             }
-
-            ColumnDescriptor::Component(desc) => Ok(Self::Component {
+            AnyColumnDescriptor::Component(desc) => Ok(Self::Component {
                 entity_path: desc.entity_path.clone(),
                 component_name: desc.component_name,
                 component_data: ComponentData::try_new(desc, column_data)?,
@@ -210,7 +221,7 @@ impl DisplayColumn {
 
     pub fn instance_count(&self, row_index: usize) -> u64 {
         match self {
-            Self::Timeline { .. } => 1,
+            Self::RowId { .. } | Self::Timeline { .. } => 1,
             Self::Component { component_data, .. } => component_data.instance_count(row_index),
         }
     }
@@ -237,6 +248,14 @@ impl DisplayColumn {
         }
 
         match self {
+            Self::RowId { row_ids } => {
+                if instance_index.is_some() {
+                    // we only ever display the row id on the summary line
+                    return;
+                }
+
+                ui.label(row_ids[row_index].to_string());
+            }
             Self::Timeline {
                 timeline,
                 time_data,
@@ -292,7 +311,7 @@ impl DisplayColumn {
                 let timestamp = time_data.get(row_index)?;
                 TimeInt::try_from(*timestamp).ok()
             }
-            Self::Component { .. } => None,
+            Self::RowId { .. } | Self::Component { .. } => None,
         }
     }
 }
@@ -309,21 +328,21 @@ impl DisplayRecordBatch {
     /// The columns in the record batch must match the selected columns. This is guaranteed by
     /// `re_datastore`.
     pub fn try_new(
-        row_data: &[ArrowArrayRef],
-        selected_columns: &[ColumnDescriptor],
+        data: impl Iterator<Item = (AnyColumnDescriptor, ArrowArrayRef)>,
     ) -> Result<Self, DisplayRecordBatchError> {
-        let num_rows = row_data.first().map(|arr| arr.len()).unwrap_or(0);
+        let mut num_rows = None;
 
-        let columns: Result<Vec<_>, _> = selected_columns
-            .iter()
-            .zip(row_data)
+        let columns: Result<Vec<_>, _> = data
             .map(|(column_descriptor, column_data)| {
-                DisplayColumn::try_new(column_descriptor, column_data)
+                if num_rows.is_none() {
+                    num_rows = Some(column_data.len());
+                }
+                DisplayColumn::try_new(&column_descriptor, &column_data)
             })
             .collect();
 
         Ok(Self {
-            num_rows,
+            num_rows: num_rows.unwrap_or(0),
             columns: columns?,
         })
     }