vortex_layout/
display.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use termtree::Tree;
5use vortex_error::VortexResult;
6
7use crate::LayoutRef;
8
9/// Display wrapper for layout tree visualization
10pub struct DisplayLayoutTree {
11    layout: LayoutRef,
12    verbose: bool,
13}
14
15impl DisplayLayoutTree {
16    pub fn new(layout: LayoutRef, verbose: bool) -> Self {
17        Self { layout, verbose }
18    }
19}
20
21impl std::fmt::Display for DisplayLayoutTree {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        fn make_tree(layout: LayoutRef, verbose: bool) -> VortexResult<Tree<String>> {
24            // Build the node label with encoding, dtype, and metadata
25            let mut node_parts = vec![
26                format!("{}", layout.encoding()),
27                format!("dtype: {}", layout.dtype()),
28            ];
29
30            // Add child count if there are children
31            let nchildren = layout.nchildren();
32            if nchildren > 0 {
33                node_parts.push(format!("children: {}", nchildren));
34            }
35
36            // Add metadata information if verbose mode is enabled
37            if verbose {
38                let metadata = layout.metadata();
39                if !metadata.is_empty() {
40                    node_parts.push(format!("metadata: {} bytes", metadata.len()));
41                }
42
43                // Add segment IDs
44                let segment_ids = layout.segment_ids();
45                if !segment_ids.is_empty() {
46                    node_parts.push(format!(
47                        "segments: [{}]",
48                        segment_ids
49                            .iter()
50                            .map(|s| s.to_string())
51                            .collect::<Vec<_>>()
52                            .join(", ")
53                    ));
54                }
55
56                // Add row count
57                node_parts.push(format!("rows: {}", layout.row_count()));
58            }
59
60            let node_name = node_parts.join(", ");
61
62            // Get children and child names directly from the layout (not loading arrays)
63            let children = layout.children()?;
64            let child_names: Vec<_> = layout.child_names().collect();
65
66            // Build child trees
67            let child_trees: VortexResult<Vec<Tree<String>>> =
68                if !children.is_empty() && child_names.len() == children.len() {
69                    // If we have names for all children, use them
70                    children
71                        .into_iter()
72                        .zip(child_names.iter())
73                        .map(|(child, name)| {
74                            let child_tree = make_tree(child, verbose)?;
75                            Ok(Tree::new(format!("{}: {}", name, child_tree.root))
76                                .with_leaves(child_tree.leaves))
77                        })
78                        .collect()
79                } else if !children.is_empty() {
80                    // No names available, just show children
81                    children
82                        .into_iter()
83                        .map(|c| make_tree(c, verbose))
84                        .collect()
85                } else {
86                    // Leaf node - no children
87                    Ok(Vec::new())
88                };
89
90            Ok(Tree::new(node_name).with_leaves(child_trees?))
91        }
92
93        match make_tree(self.layout.clone(), self.verbose) {
94            Ok(tree) => write!(f, "{}", tree),
95            Err(e) => write!(f, "Error building layout tree: {}", e),
96        }
97    }
98}