vortex_array/array/
mod.rs

1mod canonical;
2mod compute;
3mod convert;
4mod implementation;
5mod statistics;
6mod validity;
7mod variants;
8mod visitor;
9
10use std::any::Any;
11use std::fmt::{Debug, Display, Formatter};
12use std::sync::Arc;
13
14pub use canonical::*;
15pub use compute::*;
16pub use convert::*;
17pub use implementation::*;
18pub use statistics::*;
19pub use validity::*;
20pub use variants::*;
21pub use visitor::*;
22use vortex_dtype::DType;
23use vortex_error::{VortexExpect, VortexResult};
24use vortex_mask::Mask;
25
26use crate::arrays::{
27    BoolEncoding, ExtensionEncoding, ListEncoding, NullEncoding, PrimitiveEncoding, StructEncoding,
28    VarBinEncoding, VarBinViewEncoding,
29};
30use crate::builders::ArrayBuilder;
31use crate::compute::{ComputeFn, KernelRef};
32use crate::stats::StatsSetRef;
33use crate::vtable::{EncodingVTable, VTableRef};
34use crate::{Canonical, EncodingId};
35
36/// The base trait for all Vortex arrays.
37///
38/// Users should invoke functions on this trait. Implementations should implement the corresponding
39/// function on the `_Impl` traits, e.g. [`ArrayValidityImpl`]. The functions here dispatch to the
40/// implementations, while validating pre- and post-conditions.
41pub trait Array: Send + Sync + Debug + ArrayStatistics + ArrayVariants + ArrayVisitor {
42    /// Returns the array as a reference to a generic [`Any`] trait object.
43    fn as_any(&self) -> &dyn Any;
44
45    /// Returns the array as an [`Arc`] reference to a generic [`Any`] trait object.
46    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync>;
47
48    /// Returns the array as an [`ArrayRef`].
49    fn to_array(&self) -> ArrayRef;
50
51    /// Converts the array into an [`ArrayRef`].
52    fn into_array(self) -> ArrayRef
53    where
54        Self: Sized;
55
56    /// Returns the length of the array.
57    fn len(&self) -> usize;
58
59    /// Returns whether the array is empty (has zero rows).
60    fn is_empty(&self) -> bool {
61        self.len() == 0
62    }
63
64    /// Returns the logical Vortex [`DType`] of the array.
65    fn dtype(&self) -> &DType;
66
67    /// Returns the encoding of the array.
68    fn encoding(&self) -> EncodingId;
69
70    /// Returns the encoding VTable.
71    fn vtable(&self) -> VTableRef;
72
73    /// Attempts to find a kernel for the given compute invocation.
74    fn find_kernel(&self, compute_fn: &dyn ComputeFn) -> Option<KernelRef>;
75
76    /// Returns whether the array is of the given encoding.
77    fn is_encoding(&self, encoding: EncodingId) -> bool {
78        self.encoding() == encoding
79    }
80
81    /// Returns whether this array is an arrow encoding.
82    // TODO(ngates): this shouldn't live here.
83    fn is_arrow(&self) -> bool {
84        self.is_encoding(NullEncoding.id())
85            || self.is_encoding(BoolEncoding.id())
86            || self.is_encoding(PrimitiveEncoding.id())
87            || self.is_encoding(VarBinEncoding.id())
88            || self.is_encoding(VarBinViewEncoding.id())
89    }
90
91    /// Whether the array is of a canonical encoding.
92    // TODO(ngates): this shouldn't live here.
93    fn is_canonical(&self) -> bool {
94        self.is_encoding(NullEncoding.id())
95            || self.is_encoding(BoolEncoding.id())
96            || self.is_encoding(PrimitiveEncoding.id())
97            || self.is_encoding(StructEncoding.id())
98            || self.is_encoding(ListEncoding.id())
99            || self.is_encoding(VarBinViewEncoding.id())
100            || self.is_encoding(ExtensionEncoding.id())
101    }
102
103    /// Returns whether the item at `index` is valid.
104    fn is_valid(&self, index: usize) -> VortexResult<bool>;
105
106    /// Returns whether the item at `index` is invalid.
107    fn is_invalid(&self, index: usize) -> VortexResult<bool>;
108
109    /// Returns whether all items in the array are valid.
110    ///
111    /// This is usually cheaper than computing a precise `valid_count`.
112    fn all_valid(&self) -> VortexResult<bool>;
113
114    /// Returns whether the array is all invalid.
115    ///
116    /// This is usually cheaper than computing a precise `invalid_count`.
117    fn all_invalid(&self) -> VortexResult<bool>;
118
119    /// Returns the number of valid elements in the array.
120    fn valid_count(&self) -> VortexResult<usize>;
121
122    /// Returns the number of invalid elements in the array.
123    fn invalid_count(&self) -> VortexResult<usize>;
124
125    /// Returns the canonical validity mask for the array.
126    fn validity_mask(&self) -> VortexResult<Mask>;
127
128    /// Returns the canonical representation of the array.
129    fn to_canonical(&self) -> VortexResult<Canonical>;
130
131    /// Writes the array into the canonical builder.
132    ///
133    /// The [`DType`] of the builder must match that of the array.
134    fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()>;
135
136    /// Returns the statistics of the array.
137    // TODO(ngates): change how this works. It's weird.
138    fn statistics(&self) -> StatsSetRef<'_>;
139
140    /// Replaces the children of the array with the given array references.
141    fn with_children(&self, children: &[ArrayRef]) -> VortexResult<ArrayRef>;
142}
143
144impl Array for Arc<dyn Array> {
145    fn as_any(&self) -> &dyn Any {
146        self.as_ref().as_any()
147    }
148
149    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
150        self
151    }
152
153    fn to_array(&self) -> ArrayRef {
154        self.clone()
155    }
156
157    fn into_array(self) -> ArrayRef {
158        self
159    }
160
161    fn len(&self) -> usize {
162        self.as_ref().len()
163    }
164
165    fn dtype(&self) -> &DType {
166        self.as_ref().dtype()
167    }
168
169    fn encoding(&self) -> EncodingId {
170        self.as_ref().encoding()
171    }
172
173    fn vtable(&self) -> VTableRef {
174        self.as_ref().vtable()
175    }
176
177    fn find_kernel(&self, compute_fn: &dyn ComputeFn) -> Option<KernelRef> {
178        self.as_ref().find_kernel(compute_fn)
179    }
180
181    fn is_valid(&self, index: usize) -> VortexResult<bool> {
182        self.as_ref().is_valid(index)
183    }
184
185    fn is_invalid(&self, index: usize) -> VortexResult<bool> {
186        self.as_ref().is_invalid(index)
187    }
188
189    fn all_valid(&self) -> VortexResult<bool> {
190        self.as_ref().all_valid()
191    }
192
193    fn all_invalid(&self) -> VortexResult<bool> {
194        self.as_ref().all_invalid()
195    }
196
197    fn valid_count(&self) -> VortexResult<usize> {
198        self.as_ref().valid_count()
199    }
200
201    fn invalid_count(&self) -> VortexResult<usize> {
202        self.as_ref().invalid_count()
203    }
204
205    fn validity_mask(&self) -> VortexResult<Mask> {
206        self.as_ref().validity_mask()
207    }
208
209    fn to_canonical(&self) -> VortexResult<Canonical> {
210        self.as_ref().to_canonical()
211    }
212
213    fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
214        self.as_ref().append_to_builder(builder)
215    }
216
217    fn statistics(&self) -> StatsSetRef<'_> {
218        self.as_ref().statistics()
219    }
220
221    fn with_children(&self, children: &[ArrayRef]) -> VortexResult<ArrayRef> {
222        self.as_ref().with_children(children)
223    }
224}
225
226/// A reference counted pointer to a dynamic [`Array`] trait object.
227pub type ArrayRef = Arc<dyn Array>;
228
229impl ToOwned for dyn Array {
230    type Owned = ArrayRef;
231
232    fn to_owned(&self) -> Self::Owned {
233        self.to_array()
234    }
235}
236
237impl<A: Array + Clone + 'static> TryFromArrayRef for A {
238    fn try_from_array(array: ArrayRef) -> Result<Self, ArrayRef> {
239        let fallback = array.clone();
240        if let Ok(array) = array.as_any_arc().downcast::<A>() {
241            // manually drop the fallback value so `Arc::unwrap_or_clone` doesn't always have to clone
242            drop(fallback);
243            Ok(Arc::unwrap_or_clone(array))
244        } else {
245            Err(fallback)
246        }
247    }
248}
249
250impl<A: Array + Clone + 'static> TryFromArrayRef for Arc<A> {
251    fn try_from_array(array: ArrayRef) -> Result<Self, ArrayRef> {
252        let fallback = array.clone();
253        array.as_any_arc().downcast::<A>().map_err(|_| fallback)
254    }
255}
256
257pub trait ArrayExt: Array {
258    /// Returns the array downcast to the given `A`.
259    fn as_<A: Array + 'static>(&self) -> &A {
260        self.as_any()
261            .downcast_ref::<A>()
262            .vortex_expect("Failed to downcast")
263    }
264
265    /// Returns the array downcast to the given `A`.
266    fn as_opt<A: Array + 'static>(&self) -> Option<&A> {
267        self.as_any().downcast_ref::<A>()
268    }
269
270    /// Is self an array with encoding `A`.
271    fn is<A: Array + 'static>(&self) -> bool {
272        self.as_opt::<A>().is_some()
273    }
274}
275
276impl<A: Array + ?Sized> ArrayExt for A {}
277
278impl Display for dyn Array {
279    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
280        write!(
281            f,
282            "{}({}, len={})",
283            self.encoding(),
284            self.dtype(),
285            self.len()
286        )
287    }
288}
289
290#[macro_export]
291macro_rules! try_from_array_ref {
292    ($Array:ty) => {
293        impl TryFrom<$crate::ArrayRef> for $Array {
294            type Error = vortex_error::VortexError;
295
296            fn try_from(value: $crate::ArrayRef) -> Result<Self, Self::Error> {
297                Ok(::std::sync::Arc::unwrap_or_clone(
298                    value.as_any_arc().downcast::<Self>().map_err(|_| {
299                        vortex_error::vortex_err!(
300                            "Cannot downcast to {}",
301                            std::any::type_name::<Self>()
302                        )
303                    })?,
304                ))
305            }
306        }
307    };
308}
vortex_array/array/mod.rs

vortex_array/array/
mod.rs
