vortex_array/
tree.rs

1use std::fmt::{self};
2
3use humansize::{DECIMAL, format_size};
4
5use crate::arrays::ChunkedEncoding;
6use crate::nbytes::NBytes;
7use crate::vtable::EncodingVTable;
8use crate::{Array, ArrayRef, ArrayVisitor};
9
10impl dyn Array + '_ {
11    pub fn tree_display(&self) -> impl fmt::Display {
12        TreeDisplayWrapper(self.to_array())
13    }
14}
15
16struct TreeDisplayWrapper(ArrayRef);
17
18impl fmt::Display for TreeDisplayWrapper {
19    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
20        let mut array_fmt = TreeFormatter {
21            fmt,
22            indent: "".to_string(),
23            total_size: None,
24        };
25        array_fmt.format("root", self.0.clone())
26    }
27}
28
29pub struct TreeFormatter<'a, 'b: 'a> {
30    fmt: &'a mut fmt::Formatter<'b>,
31    indent: String,
32    total_size: Option<usize>,
33}
34
35impl<'a, 'b: 'a> TreeFormatter<'a, 'b> {
36    fn format(&mut self, name: &str, array: ArrayRef) -> fmt::Result {
37        let nbytes = array.nbytes();
38        let total_size = self.total_size.unwrap_or(nbytes);
39        writeln!(
40            self,
41            "{}: {} nbytes={} ({:.2}%)",
42            name,
43            array,
44            format_size(nbytes, DECIMAL),
45            100_f64 * nbytes as f64 / total_size as f64
46        )?;
47
48        self.indent(|i| {
49            write!(i, "metadata: ")?;
50            array.metadata_fmt(i.fmt)?;
51            writeln!(i.fmt)?;
52
53            for buffer in array.buffers() {
54                writeln!(
55                    i,
56                    "buffer (align={}): {} ({:.2}%)",
57                    buffer.alignment(),
58                    format_size(buffer.len(), DECIMAL),
59                    100_f64 * buffer.len() as f64 / nbytes as f64
60                )?;
61            }
62
63            Ok(())
64        })?;
65
66        let old_total_size = self.total_size;
67        if array.is_encoding(ChunkedEncoding.id()) {
68            // Clear the total size so each chunk is treated as a new root.
69            self.total_size = None
70        } else {
71            self.total_size = Some(nbytes);
72        }
73
74        self.indent(|i| {
75            for (name, child) in array
76                .children_names()
77                .into_iter()
78                .zip(array.children().into_iter())
79            {
80                i.format(&name, child)?;
81            }
82            Ok(())
83        })?;
84
85        self.total_size = old_total_size;
86        Ok(())
87    }
88
89    fn indent<F>(&mut self, indented: F) -> fmt::Result
90    where
91        F: FnOnce(&mut TreeFormatter) -> fmt::Result,
92    {
93        let original_ident = self.indent.clone();
94        self.indent += "  ";
95        let res = indented(self);
96        self.indent = original_ident;
97        res
98    }
99
100    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> fmt::Result {
101        write!(self.fmt, "{}{}", self.indent, fmt)
102    }
103}
vortex_array/tree.rs

vortex_array/
tree.rs
