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

vortex_tui/
tree.rs
