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 vortex_utils::tree::TreeDisplayAdapter;
7use vortex_utils::tree::write_indented_tree;
8
9use crate::ArrayRef;
10use crate::display::extractor::IndentedFormatter;
11use crate::display::extractor::TreeContext;
12use crate::display::extractor::TreeExtractor;
13use crate::display::extractors::BufferExtractor;
14use crate::display::extractors::EncodingSummaryExtractor;
15use crate::display::extractors::MetadataExtractor;
16use crate::display::extractors::NbytesExtractor;
17use crate::display::extractors::StatsExtractor;
18
19/// Composable tree display builder.
20///
21/// Use `tree_display()` for the default display with all built-in extractors,
22/// or `tree_display_builder()` to start with a blank slate and compose your own:
23///
24/// ```
25/// # use vortex_array::IntoArray;
26/// # use vortex_buffer::buffer;
27/// use vortex_array::display::{EncodingSummaryExtractor, NbytesExtractor, MetadataExtractor, BufferExtractor};
28///
29/// let array = buffer![0_i16, 1, 2, 3, 4].into_array();
30///
31/// // Default: all built-in extractors
32/// let full = array.tree_display();
33///
34/// // Custom: pick only what you need
35/// let custom = array.tree_display_builder()
36///     .with(EncodingSummaryExtractor)
37///     .with(NbytesExtractor)
38///     .with(MetadataExtractor);
39/// ```
40pub struct TreeDisplay {
41    array: ArrayRef,
42    extractors: Vec<Box<dyn TreeExtractor<ArrayRef, TreeContext>>>,
43}
44
45impl TreeDisplay {
46    /// Create a new tree display for the given array with no extractors.
47    ///
48    /// With no extractors, only node names and the tree structure are shown.
49    /// Use [`Self::default_display`] for the standard set of all built-in extractors.
50    pub fn new(array: ArrayRef) -> Self {
51        Self {
52            array,
53            extractors: Vec::new(),
54        }
55    }
56
57    /// Create a tree display with all built-in extractors: encoding summary, nbytes, stats,
58    /// metadata, and buffers.
59    pub fn default_display(array: ArrayRef) -> Self {
60        Self::new(array)
61            .with(EncodingSummaryExtractor)
62            .with(NbytesExtractor)
63            .with(StatsExtractor)
64            .with(MetadataExtractor)
65            .with(BufferExtractor { show_percent: true })
66    }
67
68    /// Add an extractor to the display pipeline.
69    pub fn with<E: TreeExtractor<ArrayRef, TreeContext> + 'static>(mut self, extractor: E) -> Self {
70        self.extractors.push(Box::new(extractor));
71        self
72    }
73
74    /// Add a pre-boxed extractor to the display pipeline.
75    pub fn with_boxed(mut self, extractor: Box<dyn TreeExtractor<ArrayRef, TreeContext>>) -> Self {
76        self.extractors.push(extractor);
77        self
78    }
79}
80
81impl TreeDisplayAdapter for TreeDisplay {
82    type Context = TreeContext;
83    type Node = ArrayRef;
84
85    fn write_node(
86        &self,
87        array: &ArrayRef,
88        ctx: &TreeContext,
89        f: &mut fmt::Formatter<'_>,
90    ) -> fmt::Result {
91        for extractor in &self.extractors {
92            extractor.write_header(array, ctx, f)?;
93        }
94        Ok(())
95    }
96
97    fn write_details(
98        &self,
99        array: &ArrayRef,
100        ctx: &TreeContext,
101        f: &mut IndentedFormatter<'_, '_>,
102    ) -> fmt::Result {
103        for extractor in &self.extractors {
104            extractor.write_details(array, ctx, f)?;
105        }
106        Ok(())
107    }
108
109    fn visit_children(
110        &self,
111        array: &ArrayRef,
112        visit: &mut dyn FnMut(&str, &ArrayRef, bool) -> fmt::Result,
113    ) -> fmt::Result {
114        let mut children = array
115            .children_names()
116            .into_iter()
117            .zip(array.children())
118            .peekable();
119        while let Some((child_name, child)) = children.next() {
120            let is_last = children.peek().is_none();
121            visit(&child_name, &child, is_last)?;
122        }
123        Ok(())
124    }
125}
126
127impl fmt::Display for TreeDisplay {
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        let mut ctx = TreeContext::new();
130        write_indented_tree(self, "root", &self.array, &mut ctx, f)
131    }
132}
vortex_array/display/tree_display.rs

vortex_array/display/
tree_display.rs
