Skip to main content

vortex_array/display/
extractor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4pub use vortex_utils::tree::IndentedFormatter;
5use vortex_utils::tree::TreeDisplayContext;
6pub use vortex_utils::tree::TreeDisplayExtractor as TreeExtractor;
7
8use crate::ArrayRef;
9use crate::arrays::Chunked;
10
11/// Context threaded through tree traversal for percentage calculations etc.
12pub struct TreeContext {
13    /// Stack of ancestor nbytes values. `None` entries reset the percentage root
14    /// (e.g. for chunked arrays where each chunk is its own root).
15    pub(crate) ancestor_sizes: Vec<Option<u64>>,
16}
17
18impl TreeContext {
19    pub(crate) fn new() -> Self {
20        Self {
21            ancestor_sizes: Vec::new(),
22        }
23    }
24
25    /// The total size used as the denominator for percentage calculations.
26    /// Returns `None` if there is no ancestor (i.e., this node is the root or
27    /// a chunk boundary reset the percentage root).
28    pub fn parent_total_size(&self) -> Option<u64> {
29        self.ancestor_sizes.last().cloned().flatten()
30    }
31}
32
33impl TreeDisplayContext<ArrayRef> for TreeContext {
34    fn push_parent(&mut self, parent: &ArrayRef) {
35        self.ancestor_sizes.push(if parent.is::<Chunked>() {
36            None
37        } else {
38            Some(parent.nbytes())
39        });
40    }
41
42    fn pop_parent(&mut self, _parent: &ArrayRef) {
43        self.ancestor_sizes.pop();
44    }
45}