vortex_array/vtable/
visitor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use crate::ArrayBufferVisitor;
5use crate::ArrayChildVisitor;
6use crate::ArrayChildVisitorUnnamed;
7use crate::ArrayRef;
8use crate::IntoArray;
9use crate::arrays::ConstantArray;
10use crate::buffer::BufferHandle;
11use crate::validity::Validity;
12use crate::vtable::VTable;
13
14/// Returns the validity as a child array if it produces one.
15///
16/// - `NonNullable` and `AllValid` produce no child (returns `None`)
17/// - `AllInvalid` produces a `ConstantArray` of `false` values
18/// - `Array` returns the validity array
19#[inline]
20pub fn validity_to_child(validity: &Validity, len: usize) -> Option<ArrayRef> {
21    match validity {
22        Validity::NonNullable | Validity::AllValid => None,
23        Validity::AllInvalid => Some(ConstantArray::new(false, len).into_array()),
24        Validity::Array(array) => Some(array.clone()),
25    }
26}
27
28/// Returns 1 if validity produces a child, 0 otherwise.
29#[inline]
30pub fn validity_nchildren(validity: &Validity) -> usize {
31    match validity {
32        Validity::NonNullable | Validity::AllValid => 0,
33        Validity::AllInvalid | Validity::Array(_) => 1,
34    }
35}
36
37pub trait VisitorVTable<V: VTable> {
38    /// Visit the buffers of the array.
39    fn visit_buffers(array: &V::Array, visitor: &mut dyn ArrayBufferVisitor);
40
41    /// Count the number of buffers in the array.
42    fn nbuffers(array: &V::Array) -> usize {
43        struct NBuffers(usize);
44
45        impl ArrayBufferVisitor for NBuffers {
46            fn visit_buffer_handle(&mut self, _name: &str, _handle: &BufferHandle) {
47                self.0 += 1;
48            }
49        }
50
51        let mut visitor = NBuffers(0);
52        <V::VisitorVTable as VisitorVTable<V>>::visit_buffers(array, &mut visitor);
53        visitor.0
54    }
55
56    /// Return the names of the buffers in the array.
57    fn buffer_names(array: &V::Array) -> Vec<String> {
58        struct BufferNames(Vec<String>);
59
60        impl ArrayBufferVisitor for BufferNames {
61            fn visit_buffer_handle(&mut self, name: &str, _handle: &BufferHandle) {
62                self.0.push(name.to_string());
63            }
64        }
65
66        let mut visitor = BufferNames(Vec::new());
67        <V::VisitorVTable as VisitorVTable<V>>::visit_buffers(array, &mut visitor);
68        visitor.0
69    }
70
71    /// Visit the children of the array.
72    fn visit_children(array: &V::Array, visitor: &mut dyn ArrayChildVisitor);
73
74    /// Visit the children of the array without names.
75    ///
76    /// This is more efficient than [`Self::visit_children`] when you don't need the
77    /// child names (e.g., for counting or accessing by index). The default
78    /// implementation wraps the named visitor, but array types can override
79    /// this to avoid allocating names.
80    fn visit_children_unnamed(array: &V::Array, visitor: &mut dyn ArrayChildVisitorUnnamed) {
81        struct UnnamedWrapper<'a>(&'a mut dyn ArrayChildVisitorUnnamed);
82
83        impl ArrayChildVisitor for UnnamedWrapper<'_> {
84            fn visit_child(&mut self, _name: &str, array: &ArrayRef) {
85                self.0.visit_child(array);
86            }
87        }
88
89        <V::VisitorVTable as VisitorVTable<V>>::visit_children(array, &mut UnnamedWrapper(visitor));
90    }
91
92    /// Count the number of children in the array.
93    fn nchildren(array: &V::Array) -> usize {
94        struct NChildren(usize);
95
96        impl ArrayChildVisitorUnnamed for NChildren {
97            fn visit_child(&mut self, _array: &ArrayRef) {
98                self.0 += 1;
99            }
100        }
101
102        let mut visitor = NChildren(0);
103        <V::VisitorVTable as VisitorVTable<V>>::visit_children_unnamed(array, &mut visitor);
104        visitor.0
105    }
106
107    /// Get the nth child of the array without allocating a Vec.
108    ///
109    /// Returns `None` if the index is out of bounds.
110    fn nth_child(array: &V::Array, idx: usize) -> Option<ArrayRef> {
111        struct NthChildVisitor {
112            target_idx: usize,
113            current_idx: usize,
114            result: Option<ArrayRef>,
115        }
116
117        impl ArrayChildVisitorUnnamed for NthChildVisitor {
118            fn visit_child(&mut self, array: &ArrayRef) {
119                if self.current_idx == self.target_idx && self.result.is_none() {
120                    self.result = Some(array.clone());
121                }
122                self.current_idx += 1;
123            }
124        }
125
126        let mut visitor = NthChildVisitor {
127            target_idx: idx,
128            current_idx: 0,
129            result: None,
130        };
131        <V::VisitorVTable as VisitorVTable<V>>::visit_children_unnamed(array, &mut visitor);
132        visitor.result
133    }
134}
vortex_array/vtable/visitor.rs

vortex_array/vtable/
visitor.rs
