vortex_array/array/
visitor.rs

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

vortex_array/array/
visitor.rs
