vortex_tui/
tree.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Print tree views of Vortex files.
5
6use std::path::Path;
7use std::path::PathBuf;
8
9use serde::Serialize;
10use vortex::array::stream::ArrayStreamExt;
11use vortex::error::VortexResult;
12use vortex::file::OpenOptionsSessionExt;
13use vortex::layout::LayoutRef;
14use vortex::session::VortexSession;
15
16/// Command-line arguments for the tree command.
17#[derive(Debug, clap::Parser)]
18pub struct TreeArgs {
19    /// Which kind of tree to display.
20    #[clap(subcommand)]
21    pub mode: TreeMode,
22}
23
24/// What kind of tree to display.
25#[derive(Debug, clap::Subcommand)]
26pub enum TreeMode {
27    /// Display the array encoding tree (loads and materializes arrays)
28    Array {
29        /// Path to the Vortex file
30        file: PathBuf,
31        /// Output as JSON
32        #[arg(long)]
33        json: bool,
34    },
35    /// Display the layout tree structure (metadata only, no array loading)
36    Layout {
37        /// Path to the Vortex file
38        file: PathBuf,
39        /// Show additional metadata information including buffer sizes (requires fetching segments)
40        #[arg(short, long)]
41        verbose: bool,
42        /// Output as JSON
43        #[arg(long)]
44        json: bool,
45    },
46}
47
48/// Layout tree node for JSON output.
49#[derive(Serialize)]
50pub struct LayoutTreeNode {
51    /// Encoding name.
52    pub encoding: String,
53    /// Data type.
54    pub dtype: String,
55    /// Number of rows.
56    pub row_count: u64,
57    /// Metadata size in bytes.
58    pub metadata_bytes: usize,
59    /// Segment IDs referenced by this layout.
60    pub segment_ids: Vec<u32>,
61    /// Child layouts.
62    pub children: Vec<LayoutTreeNodeWithName>,
63}
64
65/// Layout tree node with name for JSON output.
66#[derive(Serialize)]
67pub struct LayoutTreeNodeWithName {
68    /// Child name.
69    pub name: String,
70    /// Child node data.
71    #[serde(flatten)]
72    pub node: LayoutTreeNode,
73}
74
75/// Print tree views of a Vortex file (layout tree or array tree).
76///
77/// # Errors
78///
79/// Returns an error if the file cannot be opened or read.
80pub async fn exec_tree(session: &VortexSession, args: TreeArgs) -> VortexResult<()> {
81    match args.mode {
82        TreeMode::Array { file, json } => exec_array_tree(session, &file, json).await?,
83        TreeMode::Layout {
84            file,
85            verbose,
86            json,
87        } => exec_layout_tree(session, &file, verbose, json).await?,
88    }
89
90    Ok(())
91}
92
93async fn exec_array_tree(session: &VortexSession, file: &Path, _json: bool) -> VortexResult<()> {
94    let full = session
95        .open_options()
96        .open_path(file)
97        .await?
98        .scan()?
99        .into_array_stream()?
100        .read_all()
101        .await?;
102
103    println!("{}", full.display_tree());
104
105    Ok(())
106}
107
108async fn exec_layout_tree(
109    session: &VortexSession,
110    file: &Path,
111    verbose: bool,
112    json: bool,
113) -> VortexResult<()> {
114    let vxf = session.open_options().open_path(file).await?;
115    let footer = vxf.footer();
116
117    if json {
118        let tree = layout_to_json(footer.layout().clone())?;
119        let json_output = serde_json::to_string_pretty(&tree)
120            .map_err(|e| vortex::error::vortex_err!("Failed to serialize JSON: {e}"))?;
121        println!("{json_output}");
122    } else if verbose {
123        // In verbose mode, fetch segments to display buffer sizes.
124        let output = footer
125            .layout()
126            .display_tree_with_segments(vxf.segment_source())
127            .await?;
128        println!("{output}");
129    } else {
130        // In non-verbose mode, just display layout tree without fetching segments.
131        println!("{}", footer.layout().display_tree());
132    }
133
134    Ok(())
135}
136
137fn layout_to_json(layout: LayoutRef) -> VortexResult<LayoutTreeNode> {
138    let children = layout.children()?;
139    let child_names: Vec<_> = layout.child_names().collect();
140
141    let children_json: Vec<LayoutTreeNodeWithName> = children
142        .into_iter()
143        .zip(child_names.into_iter())
144        .map(|(child, name)| {
145            let node = layout_to_json(child)?;
146            Ok(LayoutTreeNodeWithName {
147                name: name.to_string(),
148                node,
149            })
150        })
151        .collect::<VortexResult<Vec<_>>>()?;
152
153    Ok(LayoutTreeNode {
154        encoding: layout.encoding().to_string(),
155        dtype: layout.dtype().to_string(),
156        row_count: layout.row_count(),
157        metadata_bytes: layout.metadata().len(),
158        segment_ids: layout.segment_ids().iter().map(|s| **s).collect(),
159        children: children_json,
160    })
161}
vortex_tui/tree.rs

vortex_tui/
tree.rs
