Skip to main content

vortex_array/display/
tree_display.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5
6use crate::ArrayRef;
7use crate::arrays::Chunked;
8use crate::display::extractor::IndentedFormatter;
9use crate::display::extractor::TreeContext;
10use crate::display::extractor::TreeExtractor;
11use crate::display::extractors::BufferExtractor;
12use crate::display::extractors::EncodingSummaryExtractor;
13use crate::display::extractors::MetadataExtractor;
14use crate::display::extractors::NbytesExtractor;
15use crate::display::extractors::StatsExtractor;
16
17/// Composable tree display builder.
18///
19/// Use `tree_display()` for the default display with all built-in extractors,
20/// or `tree_display_builder()` to start with a blank slate and compose your own:
21///
22/// ```
23/// # use vortex_array::IntoArray;
24/// # use vortex_buffer::buffer;
25/// use vortex_array::display::{EncodingSummaryExtractor, NbytesExtractor, MetadataExtractor, BufferExtractor};
26///
27/// let array = buffer![0_i16, 1, 2, 3, 4].into_array();
28///
29/// // Default: all built-in extractors
30/// let full = array.tree_display();
31///
32/// // Custom: pick only what you need
33/// let custom = array.tree_display_builder()
34///     .with(EncodingSummaryExtractor)
35///     .with(NbytesExtractor)
36///     .with(MetadataExtractor);
37/// ```
38pub struct TreeDisplay {
39    array: ArrayRef,
40    extractors: Vec<Box<dyn TreeExtractor>>,
41}
42
43impl TreeDisplay {
44    /// Create a new tree display for the given array with no extractors.
45    ///
46    /// With no extractors, only node names and the tree structure are shown.
47    /// Use [`Self::default_display`] for the standard set of all built-in extractors.
48    pub fn new(array: ArrayRef) -> Self {
49        Self {
50            array,
51            extractors: Vec::new(),
52        }
53    }
54
55    /// Create a tree display with all built-in extractors: encoding summary, nbytes, stats,
56    /// metadata, and buffers.
57    pub fn default_display(array: ArrayRef) -> Self {
58        Self::new(array)
59            .with(EncodingSummaryExtractor)
60            .with(NbytesExtractor)
61            .with(StatsExtractor)
62            .with(MetadataExtractor)
63            .with(BufferExtractor { show_percent: true })
64    }
65
66    /// Add an extractor to the display pipeline.
67    pub fn with<E: TreeExtractor + 'static>(mut self, extractor: E) -> Self {
68        self.extractors.push(Box::new(extractor));
69        self
70    }
71
72    /// Add a pre-boxed extractor to the display pipeline.
73    pub fn with_boxed(mut self, extractor: Box<dyn TreeExtractor>) -> Self {
74        self.extractors.push(extractor);
75        self
76    }
77
78    /// Recursively write a node and all its descendants directly to the formatter.
79    fn write_node(
80        &self,
81        name: &str,
82        array: &ArrayRef,
83        ctx: &mut TreeContext,
84        indent: &str,
85        f: &mut fmt::Formatter<'_>,
86    ) -> fmt::Result {
87        // Header line: "{indent}{name}:{annotations...}\n"
88        write!(f, "{indent}{name}:")?;
89        for extractor in &self.extractors {
90            extractor.write_header(array, ctx, f)?;
91        }
92        writeln!(f)?;
93
94        // Detail lines
95        let child_indent = format!("{indent}  ");
96        {
97            let mut indented = IndentedFormatter::new(f, &child_indent);
98            for extractor in &self.extractors {
99                extractor.write_details(array, ctx, &mut indented)?;
100            }
101        }
102
103        // Push context for children: chunked arrays reset the percentage root
104        let child_size = if array.is::<Chunked>() {
105            None
106        } else {
107            Some(array.nbytes())
108        };
109        ctx.push(child_size);
110
111        // Recurse into children
112        for (child_name, child) in array.children_names().into_iter().zip(array.children()) {
113            self.write_node(&child_name, &child, ctx, &child_indent, f)?;
114        }
115
116        ctx.pop();
117
118        Ok(())
119    }
120}
121
122impl fmt::Display for TreeDisplay {
123    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124        let mut ctx = TreeContext::new();
125        self.write_node("root", &self.array, &mut ctx, "", f)
126    }
127}