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::sync::Arc;
6
7use vortex_buffer::ByteBuffer;
8use vortex_error::VortexResult;
9
10use crate::Array;
11use crate::ArrayRef;
12use crate::buffer::BufferHandle;
13
14pub trait ArrayVisitor {
15    /// Returns the children of the array.
16    fn children(&self) -> Vec<ArrayRef>;
17
18    /// Returns the number of children of the array.
19    fn nchildren(&self) -> usize;
20
21    /// Returns the nth child of the array without allocating a Vec.
22    ///
23    /// Returns `None` if the index is out of bounds.
24    fn nth_child(&self, idx: usize) -> Option<ArrayRef>;
25
26    /// Returns the names of the children of the array.
27    fn children_names(&self) -> Vec<String>;
28
29    /// Returns the array's children with their names.
30    fn named_children(&self) -> Vec<(String, ArrayRef)>;
31
32    /// Returns the buffers of the array.
33    fn buffers(&self) -> Vec<ByteBuffer>;
34
35    /// Returns the buffer handles of the array.
36    fn buffer_handles(&self) -> Vec<BufferHandle>;
37
38    /// Returns the names of the buffers of the array.
39    fn buffer_names(&self) -> Vec<String>;
40
41    /// Returns the array's buffers with their names.
42    fn named_buffers(&self) -> Vec<(String, BufferHandle)>;
43
44    /// Returns the number of buffers of the array.
45    fn nbuffers(&self) -> usize;
46
47    /// Returns the serialized metadata of the array, or `None` if the array does not
48    /// support serialization.
49    fn metadata(&self) -> VortexResult<Option<Vec<u8>>>;
50
51    /// Formats a human-readable metadata description.
52    fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result;
53
54    /// Checks if all buffers in the array tree are host-resident.
55    ///
56    /// This will fail if any buffers of self or child arrays are GPU-resident.
57    fn is_host(&self) -> bool;
58}
59
60impl ArrayVisitor for Arc<dyn Array> {
61    fn children(&self) -> Vec<ArrayRef> {
62        self.as_ref().children()
63    }
64
65    fn nchildren(&self) -> usize {
66        self.as_ref().nchildren()
67    }
68
69    fn nth_child(&self, idx: usize) -> Option<ArrayRef> {
70        self.as_ref().nth_child(idx)
71    }
72
73    fn children_names(&self) -> Vec<String> {
74        self.as_ref().children_names()
75    }
76
77    fn named_children(&self) -> Vec<(String, ArrayRef)> {
78        self.as_ref().named_children()
79    }
80
81    fn buffers(&self) -> Vec<ByteBuffer> {
82        self.as_ref().buffers()
83    }
84
85    fn buffer_handles(&self) -> Vec<BufferHandle> {
86        self.as_ref().buffer_handles()
87    }
88
89    fn buffer_names(&self) -> Vec<String> {
90        self.as_ref().buffer_names()
91    }
92
93    fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
94        self.as_ref().named_buffers()
95    }
96
97    fn nbuffers(&self) -> usize {
98        self.as_ref().nbuffers()
99    }
100
101    fn metadata(&self) -> VortexResult<Option<Vec<u8>>> {
102        self.as_ref().metadata()
103    }
104
105    fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
106        self.as_ref().metadata_fmt(f)
107    }
108
109    fn is_host(&self) -> bool {
110        self.as_ref().is_host()
111    }
112}
113
114pub trait ArrayVisitorExt: Array {
115    /// Count the number of buffers encoded by self and all child arrays.
116    fn nbuffers_recursive(&self) -> usize {
117        self.children()
118            .iter()
119            .map(ArrayVisitorExt::nbuffers_recursive)
120            .sum::<usize>()
121            + self.nbuffers()
122    }
123
124    /// Depth-first traversal of the array and its children.
125    fn depth_first_traversal(&self) -> impl Iterator<Item = ArrayRef> {
126        /// A depth-first pre-order iterator over an Array.
127        struct ArrayChildrenIterator {
128            stack: Vec<ArrayRef>,
129        }
130
131        impl Iterator for ArrayChildrenIterator {
132            type Item = ArrayRef;
133
134            fn next(&mut self) -> Option<Self::Item> {
135                let next = self.stack.pop()?;
136                for child in next.children().into_iter().rev() {
137                    self.stack.push(child);
138                }
139                Some(next)
140            }
141        }
142
143        ArrayChildrenIterator {
144            stack: vec![self.to_array()],
145        }
146    }
147}
148
149impl<A: Array + ?Sized> ArrayVisitorExt for A {}
vortex_array/array/visitor.rs

vortex_array/array/
visitor.rs
