feat(vx): add segments command for displaying segment information

palaska · palaska · commit bcc85a61c9e1 · 2025-12-07T04:19:21.000Z
Adds a new 'vx segments' command that outputs structured JSON information
about all segments in a Vortex file, including byte offsets, row ranges,
and alignment information. The segment tree logic is shared with the TUI.

Signed-off-by: Baris Palaska &lt;barispalaska@gmail.com&gt;
diff --git a/vortex-tui/src/main.rs b/vortex-tui/src/main.rs
@@ -5,6 +5,8 @@
 mod browse;
 mod convert;
 mod inspect;
+mod segment_tree;
+mod segments;
 mod tree;
 
 use std::path::PathBuf;
@@ -21,8 +23,10 @@ use vortex::io::session::RuntimeSessionExt;
 use vortex::session::VortexSession;
 
 use crate::inspect::InspectArgs;
+use crate::segments::SegmentsArgs;
 
 #[derive(clap::Parser)]
+#[command(version)]
 struct Cli {
     #[clap(subcommand)]
     command: Commands,
@@ -38,18 +42,21 @@ enum Commands {
     Browse { file: PathBuf },
     /// Inspect Vortex file footer and metadata
     Inspect(InspectArgs),
+    /// Display segment information for a Vortex file
+    Segments(SegmentsArgs),
 }
 
 impl Commands {
     fn file_path(&self) -> &PathBuf {
         match self {
             Commands::Tree(args) => match &args.mode {
-                tree::TreeMode::Array { file } => file,
+                tree::TreeMode::Array { file, .. } => file,
                 tree::TreeMode::Layout { file, .. } => file,
             },
             Commands::Browse { file } => file,
             Commands::Convert(flags) => &flags.file,
             Commands::Inspect(args) => &args.file,
+            Commands::Segments(args) => &args.file,
         }
     }
 }
@@ -81,6 +88,7 @@ async fn main() -> anyhow::Result<()> {
         Commands::Convert(flags) => convert::exec_convert(flags).await?,
         Commands::Browse { file } => exec_tui(file).await?,
         Commands::Inspect(args) => inspect::exec_inspect(args).await?,
+        Commands::Segments(args) => segments::exec_segments(args).await?,
     };
 
     Ok(())
diff --git a/vortex-tui/src/segment_tree.rs b/vortex-tui/src/segment_tree.rs
@@ -0,0 +1,141 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Shared segment tree collection logic used by both the TUI browse view and the CLI segments command.
+
+use std::sync::Arc;
+
+use vortex::dtype::FieldName;
+use vortex::error::VortexResult;
+use vortex::file::SegmentSpec;
+use vortex::layout::Layout;
+use vortex::layout::LayoutChildType;
+use vortex::utils::aliases::hash_map::HashMap;
+
+/// Information about a single segment for display purposes.
+pub struct SegmentDisplay {
+    /// Name of the segment (e.g., "data", "[0]", "zones")
+    pub name: FieldName,
+    /// The underlying segment specification
+    pub spec: SegmentSpec,
+    /// Row offset within the file
+    pub row_offset: u64,
+    /// Number of rows in this segment
+    pub row_count: u64,
+}
+
+/// A tree of segments organized by field name.
+pub struct SegmentTree {
+    /// Map from field name to list of segments for that field
+    pub segments: HashMap<FieldName, Vec<SegmentDisplay>>,
+    /// Ordered list of field names (columns) in display order
+    pub segment_ordering: Vec<FieldName>,
+}
+
+/// Collect segment tree from a layout and segment map.
+pub fn collect_segment_tree(
+    root_layout: &dyn Layout,
+    segments: &Arc<[SegmentSpec]>,
+) -> SegmentTree {
+    let mut tree = SegmentTree {
+        segments: HashMap::new(),
+        segment_ordering: Vec::new(),
+    };
+    // Ignore errors during traversal - we want to collect as much as possible
+    drop(segments_by_name_impl(
+        root_layout,
+        None,
+        None,
+        Some(0),
+        segments,
+        &mut tree,
+    ));
+    tree
+}
+
+fn segments_by_name_impl(
+    root: &dyn Layout,
+    group_name: Option<FieldName>,
+    name: Option<FieldName>,
+    row_offset: Option<u64>,
+    segments: &Arc<[SegmentSpec]>,
+    segment_tree: &mut SegmentTree,
+) -> VortexResult<()> {
+    // Recurse into children
+    for (child, child_type) in root.children()?.into_iter().zip(root.child_types()) {
+        match child_type {
+            LayoutChildType::Transparent(sub_name) => segments_by_name_impl(
+                child.as_ref(),
+                group_name.clone(),
+                Some(
+                    name.as_ref()
+                        .map(|n| format!("{n}.{sub_name}").into())
+                        .unwrap_or_else(|| sub_name.into()),
+                ),
+                row_offset,
+                segments,
+                segment_tree,
+            )?,
+            LayoutChildType::Auxiliary(aux_name) => segments_by_name_impl(
+                child.as_ref(),
+                group_name.clone(),
+                Some(
+                    name.as_ref()
+                        .map(|n| format!("{n}.{aux_name}").into())
+                        .unwrap_or_else(|| aux_name.into()),
+                ),
+                Some(0),
+                segments,
+                segment_tree,
+            )?,
+            LayoutChildType::Chunk((idx, chunk_row_offset)) => segments_by_name_impl(
+                child.as_ref(),
+                group_name.clone(),
+                Some(
+                    name.as_ref()
+                        .map(|n| format!("{n}.[{idx}]"))
+                        .unwrap_or_else(|| format!("[{idx}]"))
+                        .into(),
+                ),
+                // Compute absolute row offset.
+                Some(chunk_row_offset + row_offset.unwrap_or(0)),
+                segments,
+                segment_tree,
+            )?,
+            LayoutChildType::Field(field_name) => {
+                // Step into a new group name
+                let new_group_name = group_name
+                    .as_ref()
+                    .map(|n| format!("{n}.{field_name}").into())
+                    .unwrap_or_else(|| field_name);
+                segment_tree.segment_ordering.push(new_group_name.clone());
+
+                segments_by_name_impl(
+                    child.as_ref(),
+                    Some(new_group_name),
+                    None,
+                    row_offset,
+                    segments,
+                    segment_tree,
+                )?
+            }
+        }
+    }
+
+    let current_segments = segment_tree
+        .segments
+        .entry(group_name.unwrap_or_else(|| FieldName::from("root")))
+        .or_default();
+
+    for segment_id in root.segment_ids() {
+        let segment_spec = segments[*segment_id as usize].clone();
+        current_segments.push(SegmentDisplay {
+            name: name.clone().unwrap_or_else(|| "<unnamed>".into()),
+            spec: segment_spec,
+            row_count: root.row_count(),
+            row_offset: row_offset.unwrap_or(0),
+        })
+    }
+
+    Ok(())
+}
diff --git a/vortex-tui/src/segments.rs b/vortex-tui/src/segments.rs
@@ -0,0 +1,105 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::path::PathBuf;
+
+use serde::Serialize;
+use vortex::error::VortexResult;
+use vortex::file::OpenOptionsSessionExt;
+
+use crate::SESSION;
+use crate::segment_tree::collect_segment_tree;
+
+#[derive(Debug, clap::Parser)]
+pub struct SegmentsArgs {
+    /// Path to the Vortex file
+    pub file: PathBuf,
+}
+
+#[derive(Serialize)]
+struct SegmentsOutput {
+    /// Columns in display order
+    columns: Vec<ColumnInfo>,
+}
+
+#[derive(Serialize)]
+struct ColumnInfo {
+    /// Field name (column header)
+    name: String,
+    /// Segments within this column
+    segments: Vec<SegmentInfo>,
+}
+
+#[derive(Serialize)]
+struct SegmentInfo {
+    /// Segment name (e.g., "[0]", "data", etc.)
+    name: String,
+    /// Row range start
+    row_offset: u64,
+    /// Number of rows
+    row_count: u64,
+    /// Byte offset in file
+    byte_offset: u64,
+    /// Length in bytes
+    byte_length: u32,
+    /// Alignment requirement
+    alignment: usize,
+    /// Gap from previous segment end
+    byte_gap: u64,
+}
+
+pub async fn exec_segments(args: SegmentsArgs) -> VortexResult<()> {
+    let vxf = SESSION.open_options().open(args.file).await?;
+
+    let footer = vxf.footer();
+    let mut segment_tree = collect_segment_tree(footer.layout().as_ref(), footer.segment_map());
+
+    // Convert to output format
+    let columns: Vec<ColumnInfo> = segment_tree
+        .segment_ordering
+        .iter()
+        .filter_map(|name| {
+            let mut segments = segment_tree.segments.remove(name)?;
+
+            // Sort by byte offset
+            segments.sort_by(|a, b| a.spec.offset.cmp(&b.spec.offset));
+
+            // Convert to output format, computing byte gaps
+            let mut current_offset = 0u64;
+            let segment_infos: Vec<SegmentInfo> = segments
+                .into_iter()
+                .map(|seg| {
+                    let byte_gap = if current_offset == 0 {
+                        0
+                    } else {
+                        seg.spec.offset.saturating_sub(current_offset)
+                    };
+                    current_offset = seg.spec.offset + seg.spec.length as u64;
+
+                    SegmentInfo {
+                        name: seg.name.to_string(),
+                        row_offset: seg.row_offset,
+                        row_count: seg.row_count,
+                        byte_offset: seg.spec.offset,
+                        byte_length: seg.spec.length,
+                        alignment: *seg.spec.alignment,
+                        byte_gap,
+                    }
+                })
+                .collect();
+
+            Some(ColumnInfo {
+                name: name.to_string(),
+                segments: segment_infos,
+            })
+        })
+        .collect();
+
+    let output = SegmentsOutput { columns };
+
+    let json_output = serde_json::to_string_pretty(&output)
+        .map_err(|e| vortex::error::vortex_err!("Failed to serialize JSON: {e}"))?;
+    println!("{json_output}");
+
+    Ok(())
+}
