Fix formatting of RowId/Tuid when printing ChunkStore

rerun-io · Jan 10, 2025 · eff7986 · eff7986
1 parent 9c1ed51
commit eff7986
Show file tree

Hide file tree

Showing 11 changed files with 114 additions and 25 deletions.
diff --git a/Cargo.lock b/Cargo.lock
@@ -3885,7 +3885,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4979f22fdb869068da03c9f7528f8297c6fd2606bc3a4affe42e6a823fdb8da4"
 dependencies = [
  "cfg-if",
- "windows-targets 0.48.5",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
@@ -5203,7 +5203,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0c1318b19085f08681016926435853bbf7858f9c082d0999b80550ff5d9abe15"
 dependencies = [
  "bytes",
- "heck 0.4.1",
+ "heck 0.5.0",
  "itertools 0.13.0",
  "log",
  "multimap",
@@ -5686,6 +5686,7 @@ dependencies = [
  "re_types_core",
  "serde",
  "similar-asserts",
+ "tap",
  "thiserror 1.0.65",
 ]
 
@@ -5978,6 +5979,7 @@ name = "re_format_arrow"
 version = "0.22.0-alpha.1+dev"
 dependencies = [
  "comfy-table",
+ "itertools 0.13.0",
  "re_arrow2",
  "re_tuid",
  "re_types_core",
@@ -8481,6 +8483,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "tap"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
+
 [[package]]
 name = "target-lexicon"
 version = "0.12.16"
@@ -9846,7 +9854,7 @@ version = "0.1.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb"
 dependencies = [
- "windows-sys 0.48.0",
+ "windows-sys 0.59.0",
 ]
 
 [[package]]

diff --git a/Cargo.toml b/Cargo.toml
@@ -269,6 +269,7 @@ strum_macros = "0.26"
 sublime_fuzzy = "0.7"
 syn = "2.0"
 sysinfo = { version = "0.30.1", default-features = false }
+tap = "1.0.1"
 tempfile = "3.0"
 thiserror = "1.0"
 time = { version = "0.3.36", default-features = false, features = [

diff --git a/crates/store/re_chunk/Cargo.toml b/crates/store/re_chunk/Cargo.toml
@@ -63,6 +63,7 @@ itertools.workspace = true
 nohash-hasher.workspace = true
 rand = { workspace = true, features = ["std_rng"] }
 similar-asserts.workspace = true
+tap.workspace = true
 thiserror.workspace = true
 
 # Optional dependencies:

diff --git a/crates/store/re_chunk/src/chunk.rs b/crates/store/re_chunk/src/chunk.rs
@@ -1308,11 +1308,11 @@ impl Chunk {
 impl std::fmt::Display for Chunk {
     #[inline]
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let chunk = self.to_transport().map_err(|err| {
+        let transport_chunk = self.to_transport().map_err(|err| {
             re_log::error_once!("couldn't display Chunk: {err}");
             std::fmt::Error
         })?;
-        chunk.fmt(f)
+        transport_chunk.fmt(f)
     }
 }
 

diff --git a/crates/store/re_chunk/src/transport.rs b/crates/store/re_chunk/src/transport.rs
@@ -15,6 +15,7 @@ use nohash_hasher::IntMap;
 use re_byte_size::SizeBytes as _;
 use re_log_types::{EntityPath, Timeline};
 use re_types_core::{Component as _, ComponentDescriptor, Loggable as _};
+use tap::Tap as _;
 
 use crate::{chunk::ChunkComponents, Chunk, ChunkError, ChunkId, ChunkResult, RowId, TimeColumn};
 
@@ -436,7 +437,15 @@ impl Chunk {
                     row_ids.data_type().clone(),
                     false,
                 )
-                .with_metadata(TransportChunk::field_metadata_control_column()),
+                .with_metadata(
+                    TransportChunk::field_metadata_control_column().tap_mut(|metadata| {
+                        // This ensures the RowId/Tuid is formatted correctly
+                        metadata.insert(
+                            "ARROW:extension:name".to_owned(),
+                            re_tuid::Tuid::ARROW_EXTENSION_NAME.to_owned(),
+                        );
+                    }),
+                ),
             );
             columns.push(row_ids.clone().boxed());
         }

diff --git a/crates/store/re_chunk_store/tests/formatting.rs b/crates/store/re_chunk_store/tests/formatting.rs
@@ -0,0 +1,46 @@
+use std::sync::Arc;
+
+use re_chunk::{Chunk, RowId};
+use re_chunk_store::ChunkStore;
+use re_log_types::{
+    build_frame_nr, build_log_time,
+    example_components::{MyColor, MyIndex},
+    EntityPath, Time,
+};
+
+/// Ensure that `ChunkStore::to_string()` is nice and readable.
+#[test]
+fn format_chunk_store() -> anyhow::Result<()> {
+    re_log::setup_logging();
+
+    let mut store = ChunkStore::new(
+        re_log_types::StoreId::from_string(
+            re_log_types::StoreKind::Recording,
+            "test_id".to_owned(),
+        ),
+        Default::default(),
+    );
+
+    let entity_path = EntityPath::from("this/that");
+
+    let (indices1, colors1) = (MyIndex::from_iter(0..3), MyColor::from_iter(0..3));
+
+    let row_id = RowId::from_u128(32_033_410_000_000_000_000_000_000_123);
+
+    store.insert_chunk(&Arc::new(
+        Chunk::builder(entity_path.clone())
+            .with_component_batches(
+                row_id,
+                [
+                    build_frame_nr(1),
+                    build_log_time(Time::from_ns_since_epoch(1_736_534_622_123_456_789)),
+                ],
+                [&indices1 as _, &colors1 as _],
+            )
+            .build()?,
+    ))?;
+
+    insta::assert_snapshot!("format_chunk_store", store.to_string());
+
+    Ok(())
+}
diff --git a/crates/store/re_chunk_store/tests/snapshots/formatting__format_chunk_store.snap b/crates/store/re_chunk_store/tests/snapshots/formatting__format_chunk_store.snap
@@ -0,0 +1,35 @@
+---
+source: crates/store/re_chunk_store/tests/formatting.rs
+assertion_line: 43
+expression: store.to_string()
+snapshot_kind: text
+---
+ChunkStore {
+    id: test_id
+    config: ChunkStoreConfig { enable_changelog: true, chunk_max_bytes: 393216, chunk_max_rows: 4096, chunk_max_rows_if_unsorted: 1024 }
+    stats: {
+        num_chunks: 1
+        total_size_bytes: 1.1 KiB
+        num_rows: 1
+        num_events: 2
+    }
+    chunks: [
+        ┌──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐
+        │[3m CHUNK METADATA:                                                                                                              [0m│
+        │[3m * entity_path: "/this/that"                                                                                                  [0m│
+        │[3m * heap_size_bytes: "934"                                                                                                     [0m│
+        │[3m * id: "18196B315017D6D07289B50E31531230"                                                                                     [0m│
+        │[3m * is_sorted: ""                                                                                                              [0m│
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+        │ ┌──────────────────────────────────┬───────────────┬───────────────────────────────┬───────────────────┬───────────────────┐ │
+        │ │ RowId                            ┆ frame_nr      ┆ log_time                      ┆ example.MyColor   ┆ example.MyIndex   │ │
+        │ │ ---                              ┆ ---           ┆ ---                           ┆ ---               ┆ ---               │ │
+        │ │ type: "struct[2]"                ┆ type: "i64"   ┆ type: "timestamp(ns)"         ┆ type: "list[u32]" ┆ type: "list[u64]" │ │
+        │ │ ARROW:extension:name: "TUID"     ┆ is_sorted: "" ┆ is_sorted: ""                 ┆ kind: "data"      ┆ kind: "data"      │ │
+        │ │ kind: "control"                  ┆ kind: "time"  ┆ kind: "time"                  ┆                   ┆                   │ │
+        │ ╞══════════════════════════════════╪═══════════════╪═══════════════════════════════╪═══════════════════╪═══════════════════╡ │
+        │ │ 0000000067816A6BB4B8C1254D40007B ┆ 1             ┆ 2025-01-10 18:43:42.123456789 ┆ [0, 1, 2]         ┆ [0, 1, 2]         │ │
+        │ └──────────────────────────────────┴───────────────┴───────────────────────────────┴───────────────────┴───────────────────┘ │
+        └──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘
+    ]
+}
diff --git a/crates/store/re_format_arrow/Cargo.toml b/crates/store/re_format_arrow/Cargo.toml
@@ -21,6 +21,7 @@ all-features = true
 
 [dependencies]
 arrow2.workspace = true
+itertools.workspace = true
 re_tuid.workspace = true
 re_types_core.workspace = true # tuid serialization
 

diff --git a/crates/store/re_format_arrow/src/lib.rs b/crates/store/re_format_arrow/src/lib.rs
@@ -21,24 +21,13 @@ use re_types_core::Loggable as _;
 type CustomFormatter<'a, F> = Box<dyn Fn(&mut F, usize) -> std::fmt::Result + 'a>;
 
 fn get_custom_display<'a, F: std::fmt::Write + 'a>(
+    field: &Field,
     array: &'a dyn Array,
     null: &'static str,
 ) -> CustomFormatter<'a, F> {
-    // NOTE: If the top-level array is a list, it's probably not the type we're looking for: we're
-    // interested in the type of the array that's underneath.
-    let datatype = (|| match array.data_type().to_logical_type() {
-        DataType::List(_) => array
-            .as_any()
-            .downcast_ref::<ListArray<i32>>()?
-            .iter()
-            .next()?
-            .map(|array| array.data_type().clone()),
-        _ => Some(array.data_type().clone()),
-    })();
-
-    if let Some(DataType::Extension(name, _, _)) = datatype {
+    if let Some(extension_name) = field.metadata.get("ARROW:extension:name") {
         // TODO(#1775): This should be registered dynamically.
-        if name.as_str() == Tuid::NAME {
+        if extension_name.as_str() == Tuid::ARROW_EXTENSION_NAME {
             return Box::new(|w, index| {
                 if let Some(tuid) = parse_tuid(array, index) {
                     w.write_fmt(format_args!("{tuid}"))
@@ -280,9 +269,8 @@ where
     });
     table.set_header(header);
 
-    let displays = columns
-        .iter()
-        .map(|array| get_custom_display(&**array, "-"))
+    let displays = itertools::izip!(fields.iter(), columns.iter())
+        .map(|(field, array)| get_custom_display(field, &**array, "-"))
         .collect::<Vec<_>>();
     let num_rows = columns.first().map_or(0, |list_array| list_array.len());
 

diff --git a/crates/store/re_types_core/src/tuid.rs b/crates/store/re_types_core/src/tuid.rs
@@ -27,7 +27,7 @@ impl Loggable for Tuid {
         Self: 'a,
     {
         Err(crate::SerializationError::not_implemented(
-            Self::NAME,
+            Self::ARROW_EXTENSION_NAME,
             "TUIDs are never nullable, use `to_arrow()` instead",
         ))
     }

diff --git a/crates/utils/re_tuid/src/lib.rs b/crates/utils/re_tuid/src/lib.rs
@@ -23,7 +23,7 @@ impl Tuid {
     /// We give an actual name to [`Tuid`], and inject that name into the Arrow datatype extensions,
     /// as a hack so that we can compactly format them when printing Arrow data to the terminal.
     /// Check out `re_format_arrow` for context.
-    pub const NAME: &'static str = "rerun.datatypes.TUID";
+    pub const ARROW_EXTENSION_NAME: &'static str = "rerun.datatypes.TUID";
 }
 
 impl std::fmt::Display for Tuid {