vortex_array/
tree.rs

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

vortex_array/
tree.rs
