vortex_array/array/
visitor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Formatter;
5use std::ops::Deref;
6use std::sync::Arc;
7
8use vortex_buffer::ByteBuffer;
9use vortex_error::VortexResult;
10
11use crate::arrays::ConstantArray;
12use crate::patches::Patches;
13use crate::validity::Validity;
14use crate::{Array, ArrayRef};
15
16pub trait ArrayVisitor {
17    /// Returns the children of the array.
18    fn children(&self) -> Vec<ArrayRef>;
19
20    /// Returns the number of children of the array.
21    fn nchildren(&self) -> usize;
22
23    /// Returns the names of the children of the array.
24    fn children_names(&self) -> Vec<String>;
25
26    /// Returns the array's children with their names.
27    fn named_children(&self) -> Vec<(String, ArrayRef)>;
28
29    /// Returns the buffers of the array.
30    fn buffers(&self) -> Vec<ByteBuffer>;
31
32    /// Returns the number of buffers of the array.
33    fn nbuffers(&self) -> usize;
34
35    /// Returns the serialized metadata of the array, or `None` if the array does not
36    /// support serialization.
37    fn metadata(&self) -> VortexResult<Option<Vec<u8>>>;
38
39    /// Formats a human-readable metadata description.
40    fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result;
41}
42
43impl ArrayVisitor for Arc<dyn Array> {
44    fn children(&self) -> Vec<ArrayRef> {
45        self.as_ref().children()
46    }
47
48    fn nchildren(&self) -> usize {
49        self.as_ref().nchildren()
50    }
51
52    fn children_names(&self) -> Vec<String> {
53        self.as_ref().children_names()
54    }
55
56    fn named_children(&self) -> Vec<(String, ArrayRef)> {
57        self.as_ref().named_children()
58    }
59
60    fn buffers(&self) -> Vec<ByteBuffer> {
61        self.as_ref().buffers()
62    }
63
64    fn nbuffers(&self) -> usize {
65        self.as_ref().nbuffers()
66    }
67
68    fn metadata(&self) -> VortexResult<Option<Vec<u8>>> {
69        self.as_ref().metadata()
70    }
71
72    fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
73        self.as_ref().metadata_fmt(f)
74    }
75}
76
77pub trait ArrayVisitorExt: Array {
78    /// Count the number of buffers encoded by self and all child arrays.
79    fn nbuffers_recursive(&self) -> usize {
80        self.children()
81            .iter()
82            .map(ArrayVisitorExt::nbuffers_recursive)
83            .sum::<usize>()
84            + self.nbuffers()
85    }
86
87    /// Depth-first traversal of the array and its children.
88    fn depth_first_traversal(&self) -> impl Iterator<Item = ArrayRef> {
89        /// A depth-first pre-order iterator over an Array.
90        struct ArrayChildrenIterator {
91            stack: Vec<ArrayRef>,
92        }
93
94        impl Iterator for ArrayChildrenIterator {
95            type Item = ArrayRef;
96
97            fn next(&mut self) -> Option<Self::Item> {
98                let next = self.stack.pop()?;
99                for child in next.children().into_iter().rev() {
100                    self.stack.push(child);
101                }
102                Some(next)
103            }
104        }
105
106        ArrayChildrenIterator {
107            stack: vec![self.to_array()],
108        }
109    }
110}
111
112impl<A: Array + ?Sized> ArrayVisitorExt for A {}
113
114pub trait ArrayBufferVisitor {
115    fn visit_buffer(&mut self, buffer: &ByteBuffer);
116}
117
118pub trait ArrayChildVisitor {
119    /// Visit a child of this array.
120    fn visit_child(&mut self, _name: &str, _array: &dyn Array);
121
122    /// Utility for visiting Array validity.
123    fn visit_validity(&mut self, validity: &Validity, len: usize) {
124        if let Some(vlen) = validity.maybe_len() {
125            assert_eq!(vlen, len, "Validity length mismatch");
126        }
127
128        match validity {
129            Validity::NonNullable | Validity::AllValid => {}
130            Validity::AllInvalid => {
131                // To avoid storing metadata about validity, we store all invalid as a
132                // constant array of false values.
133                // This gives:
134                //  * is_nullable & has_validity => Validity::Array (or Validity::AllInvalid)
135                //  * is_nullable & !has_validity => Validity::AllValid
136                //  * !is_nullable => Validity::NonNullable
137                self.visit_child("validity", ConstantArray::new(false, len).deref())
138            }
139            Validity::Array(array) => {
140                self.visit_child("validity", array);
141            }
142        }
143    }
144
145    /// Utility for visiting Array patches.
146    fn visit_patches(&mut self, patches: &Patches) {
147        self.visit_child("patch_indices", patches.indices());
148        self.visit_child("patch_values", patches.values());
149    }
150}
vortex_array/array/visitor.rs

vortex_array/array/
visitor.rs
