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            100f64 * 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={}): {}",
57                    buffer.alignment(),
58                    format_size(buffer.len(), DECIMAL)
59                )?;
60            }
61
62            Ok(())
63        })?;
64
65        let old_total_size = self.total_size;
66        if array.is_encoding(ChunkedEncoding.id()) {
67            // Clear the total size so each chunk is treated as a new root.
68            self.total_size = None
69        } else {
70            self.total_size = Some(total_size);
71        }
72
73        self.indent(|i| {
74            for (name, child) in array
75                .children_names()
76                .into_iter()
77                .zip(array.children().into_iter())
78            {
79                i.format(&name, child)?;
80            }
81            Ok(())
82        })?;
83
84        self.total_size = old_total_size;
85        Ok(())
86    }
87
88    fn indent<F>(&mut self, indented: F) -> fmt::Result
89    where
90        F: FnOnce(&mut TreeFormatter) -> fmt::Result,
91    {
92        let original_ident = self.indent.clone();
93        self.indent += "  ";
94        let res = indented(self);
95        self.indent = original_ident;
96        res
97    }
98
99    fn write_fmt(&mut self, fmt: fmt::Arguments<'_>) -> fmt::Result {
100        write!(self.fmt, "{}{}", self.indent, fmt)
101    }
102}
vortex_array/tree.rs

vortex_array/
tree.rs
