Layout segments size info in TUI (#2833)

I'm looking into compression/file size and this is just thing I'm adding
as I go. The main change here is restoring size info by traversing the
segments to get all segments contained within a layout.
Top layout info:
<img width="357" alt="Screenshot 2025-03-28 at 16 50 08"
src="https://github.com/user-attachments/assets/a2566e74-9935-471f-a6d5-70526194a3b9"
/>

child level info for chunked layout
<img width="295" alt="Screenshot 2025-03-28 at 16 50 49"
src="https://github.com/user-attachments/assets/c54bb540-6414-489e-a054-cb8b14ec2e08"
/>
and struct layout:
<img width="463" alt="Screenshot 2025-03-28 at 16 51 01"
src="https://github.com/user-attachments/assets/6d04ac22-ecf3-4a1f-92c2-0c6ea2e1828c"
/>

Author: AdamGS

Cargo.lock

@@ -25,6 +25,18 @@ impl From<&str> for Field {
     }
 }
 
+impl From<Arc<str>> for Field {
+    fn from(value: Arc<str>) -> Self {
+        Self::Name(value)
+    }
+}
+
+impl From<&Arc<str>> for Field {
+    fn from(value: &Arc<str>) -> Self {
+        Self::Name(value.clone())
+    }
+}
+
 impl From<String> for Field {
     fn from(value: String) -> Self {
         Field::Name(value.into())

vortex-dtype/src/field.rs

@@ -19,6 +19,7 @@ async-trait = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
 crossterm = { workspace = true }
 futures-util = { workspace = true }
+humansize = { workspace = true }
 indicatif = { workspace = true, features = ["futures"] }
 parquet = { workspace = true, features = ["arrow", "async"] }
 ratatui = { workspace = true }

vortex-tui/Cargo.toml

@@ -7,14 +7,19 @@ use async_trait::async_trait;
 use ratatui::widgets::ListState;
 use vortex::buffer::{Alignment, ByteBuffer, ByteBufferMut};
 use vortex::dtype::DType;
-use vortex::error::{VortexExpect, VortexResult};
+use vortex::error::{VortexExpect, VortexResult, VortexUnwrap};
 use vortex::file::{Footer, SegmentSpec, VortexOpenOptions};
 use vortex::io::TokioFile;
 use vortex::stats::stats_from_bitset_bytes;
+use vortex_layout::layouts::chunked::ChunkedLayout;
+use vortex_layout::layouts::flat::FlatLayout;
+use vortex_layout::layouts::stats::StatsLayout;
 use vortex_layout::layouts::stats::stats_table::StatsTable;
+use vortex_layout::layouts::struct_::StructLayout;
 use vortex_layout::segments::{AsyncSegmentReader, SegmentId};
 use vortex_layout::{
-    CHUNKED_LAYOUT_ID, FLAT_LAYOUT_ID, Layout, LayoutVTableRef, STATS_LAYOUT_ID, STRUCT_LAYOUT_ID,
+    CHUNKED_LAYOUT_ID, FLAT_LAYOUT_ID, Layout, LayoutVTable, LayoutVTableRef, STATS_LAYOUT_ID,
+    STRUCT_LAYOUT_ID,
 };
 
 #[derive(Default, Copy, Clone, Eq, PartialEq)]
@@ -35,7 +40,6 @@ pub struct LayoutCursor {
     path: Vec<usize>,
     footer: Footer,
     layout: Layout,
-    #[allow(unused)]
     segment_map: Arc<[SegmentSpec]>,
 }
 
@@ -137,13 +141,18 @@ impl LayoutCursor {
             .unwrap_or_default()
     }
 
-    pub fn segment_size(&self) -> usize {
-        self.layout()
-            .segments()
-            .map(|id| self.segment_spec(id).length as usize)
+    pub fn total_size(&self) -> usize {
+        self.layout_segments()
+            .iter()
+            .map(|id| self.segment_spec(*id).length as usize)
             .sum()
     }
 
+    fn layout_segments(&self) -> Vec<SegmentId> {
+        let segments = collect_segment_ids(&self.layout);
+        [segments.0, segments.1].concat()
+    }
+
     /// Predicate true when the cursor is currently activated over a stats table
     pub fn is_stats_table(&self) -> bool {
         let parent = self.parent();
@@ -262,3 +271,50 @@ impl AsyncSegmentReader for SegmentReader {
         Ok(self.read_bytes_sync(range, segment.alignment))
     }
 }
+
+pub fn collect_segment_ids(root_layout: &Layout) -> (Vec<SegmentId>, Vec<SegmentId>) {
+    let mut data_segment_ids = Vec::default();
+    let mut stats_segment_ids = Vec::default();
+
+    collect_segment_ids_impl(root_layout, &mut data_segment_ids, &mut stats_segment_ids)
+        .vortex_unwrap();
+
+    (data_segment_ids, stats_segment_ids)
+}
+
+fn collect_segment_ids_impl(
+    root: &Layout,
+    data_segments: &mut Vec<SegmentId>,
+    stats_segments: &mut Vec<SegmentId>,
+) -> VortexResult<()> {
+    let layout_id = root.id();
+
+    if layout_id == StructLayout.id() {
+        let dtype = root.dtype().as_struct().vortex_expect("");
+        for child_idx in 0..dtype.fields().len() {
+            let name = dtype.field_name(child_idx)?;
+            let child_dtype = dtype.field_by_index(child_idx)?;
+            let child_layout = root.child(child_idx, child_dtype, name)?;
+            collect_segment_ids_impl(&child_layout, data_segments, stats_segments)?;
+        }
+    } else if layout_id == ChunkedLayout.id() {
+        for child_idx in 0..root.nchildren() {
+            let child_layout =
+                root.child(child_idx, root.dtype().clone(), format!("[{child_idx}]"))?;
+            collect_segment_ids_impl(&child_layout, data_segments, stats_segments)?;
+        }
+    } else if layout_id == StatsLayout.id() {
+        let data_layout = root.child(0, root.dtype().clone(), "data")?;
+        collect_segment_ids_impl(&data_layout, data_segments, stats_segments)?;
+
+        // For the stats layout, we use the stats segment accumulator
+        let stats_layout = root.child(1, root.dtype().clone(), "stats")?;
+        collect_segment_ids_impl(&stats_layout, stats_segments, &mut vec![])?;
+    } else if layout_id == FlatLayout.id() {
+        data_segments.extend(root.segments());
+    } else {
+        unreachable!()
+    };
+
+    Ok(())
+}

vortex-tui/src/browse/app.rs

@@ -1,3 +1,4 @@
+use humansize::{DECIMAL, format_size, make_format};
 use ratatui::buffer::Buffer;
 use ratatui::layout::{Constraint, Layout, Rect};
 use ratatui::style::{Color, Style, Stylize};
@@ -11,7 +12,7 @@ use vortex::error::VortexExpect;
 use vortex::expr::Identity;
 use vortex::layout::{CHUNKED_LAYOUT_ID, FLAT_LAYOUT_ID, STATS_LAYOUT_ID, STRUCT_LAYOUT_ID};
 use vortex::stats::stats_from_bitset_bytes;
-use vortex_layout::layouts::flat::FlatLayout;
+use vortex_layout::layouts::stats::StatsLayout;
 use vortex_layout::{ExprEvaluator, LayoutReaderExt, LayoutVTable, RowMask};
 
 use crate::TOKIO_RUNTIME;
@@ -41,18 +42,16 @@ pub fn render_layouts(app_state: &mut AppState, area: Rect, buf: &mut Buffer) {
 fn render_layout_header(cursor: &LayoutCursor, area: Rect, buf: &mut Buffer) {
     let layout_kind = cursor.layout().id().to_string();
     let row_count = cursor.layout().row_count();
+    let size = format_size(cursor.total_size(), DECIMAL);
 
     let mut rows = vec![
         Text::from(format!("Kind: {layout_kind}")).bold(),
         Text::from(format!("Row Count: {row_count}")).bold(),
         Text::from(format!("Schema: {}", cursor.dtype()))
             .bold()
             .green(),
-        Text::from(format!("Segments: {}", cursor.layout().nsegments())),
-        Text::from(format!(
-            "Segment data size: {} bytes",
-            cursor.segment_size()
-        )),
+        Text::from(format!("Children: {}", cursor.layout().nchildren())).bold(),
+        Text::from(format!("Segment data size: {}", size)).bold(),
     ];
 
     if cursor.encoding().id() == FLAT_LAYOUT_ID {
@@ -62,7 +61,7 @@ fn render_layout_header(cursor: &LayoutCursor, area: Rect, buf: &mut Buffer) {
         )));
     }
 
-    if cursor.encoding().id() == STATS_LAYOUT_ID {
+    if cursor.encoding().id() == StatsLayout.id() {
         // Push any columnar stats.
         if let Some(available_stats) = cursor
             .layout()
@@ -97,13 +96,11 @@ fn render_layout_header(cursor: &LayoutCursor, area: Rect, buf: &mut Buffer) {
 
 // Render the inner Array for a FlatLayout
 fn render_array(app: &AppState, area: Rect, buf: &mut Buffer, is_stats_table: bool) {
-    let reader = FlatLayout
-        .reader(
-            app.cursor.layout().clone(),
-            app.footer.ctx().clone(),
-            app.reader.clone(),
-        )
-        .vortex_expect("Failed to create FlatLayout reader");
+    let reader = app
+        .cursor
+        .layout()
+        .reader(app.reader.clone(), app.footer.ctx().clone())
+        .vortex_expect("Failed to create reader");
 
     let array = TOKIO_RUNTIME
         .block_on(reader.evaluate_expr(
@@ -172,27 +169,25 @@ fn render_array(app: &AppState, area: Rect, buf: &mut Buffer, is_stats_table: bo
 }
 
 fn render_children_list(app: &mut AppState, area: Rect, buf: &mut Buffer) {
-    let cursor = &app.cursor;
     // TODO: add selection state.
-    let layout = cursor.layout();
-    let state = &mut app.layouts_list_state;
+    let search_filter = app.search_filter.clone();
 
-    if layout.nchildren() > 0 {
-        let filter: Vec<bool> = (0..layout.nchildren())
-            .map(|idx| child_name(cursor, idx))
+    if app.cursor.layout().nchildren() > 0 {
+        let filter: Vec<bool> = (0..app.cursor.layout().nchildren())
+            .map(|idx| child_name(app, idx))
             .map(|name| {
-                if app.search_filter.is_empty() {
+                if search_filter.is_empty() {
                     true
                 } else {
-                    name.contains(&app.search_filter)
+                    name.contains(&search_filter)
                 }
             })
             .collect();
 
-        let list_items: Vec<String> = (0..layout.nchildren())
+        let list_items: Vec<String> = (0..app.cursor.layout().nchildren())
             .zip(filter.iter())
             .filter(|&(_, keep)| *keep)
-            .map(|(idx, _)| child_name(cursor, idx))
+            .map(|(idx, _)| child_name(app, idx))
             .collect();
 
         if !app.search_filter.is_empty() {
@@ -215,28 +210,31 @@ fn render_children_list(app: &mut AppState, area: Rect, buf: &mut Buffer) {
             List::new(list_items).highlight_style(Style::default().black().on_white().bold()),
             inner_area,
             buf,
-            state,
+            &mut app.layouts_list_state,
         );
     }
 }
 
-fn child_name(cursor: &LayoutCursor, nth: usize) -> String {
+fn child_name(app: &mut AppState, nth: usize) -> String {
+    let cursor = &app.cursor;
+    let formatter = make_format(DECIMAL);
     // TODO(ngates): layout visitors
     if cursor.layout().id() == STRUCT_LAYOUT_ID {
         let struct_dtype = cursor.dtype().as_struct().expect("struct dtype");
         let field_name = struct_dtype.field_name(nth).expect("field name");
         let field_dtype = struct_dtype.field_by_index(nth).expect("dtype value");
-        format!("Column {nth} - {field_name} ({field_dtype})")
+
+        let total_size = formatter(app.cursor.child(nth).total_size());
+
+        format!("Column {nth} - {field_name} ({field_dtype}) - {total_size}")
     } else if cursor.layout().id() == CHUNKED_LAYOUT_ID {
-        // 0th child of a ChunkedLayout is the chunk stats array.
-        // The rest of the chunks are child arrays
-        if cursor.layout().metadata().is_none() {
-            format!("Chunk {nth}")
-        } else if nth == (cursor.layout().nchildren() - 1) {
-            "Chunk Statistics".to_string()
-        } else {
-            format!("Chunk {}", nth)
-        }
+        let name = format!("Chunk {nth}");
+        let child_cursor = app.cursor.child(nth);
+
+        let total_size = formatter(child_cursor.total_size());
+        let row_count = child_cursor.layout().row_count();
+
+        format!("{name} - {row_count} - {total_size}")
     } else if cursor.layout().id() == FLAT_LAYOUT_ID {
         format!("Page {nth}")
     } else if cursor.layout().id() == STATS_LAYOUT_ID {