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

vortex_array/array/
visitor.rs
