vortex_array/arrays/extension/compute/
mod.rs

1mod compare;
2mod to_arrow;
3
4use vortex_error::VortexResult;
5use vortex_mask::Mask;
6use vortex_scalar::Scalar;
7
8use crate::arrays::ExtensionEncoding;
9use crate::arrays::extension::ExtensionArray;
10use crate::compute::{
11    CastFn, CompareFn, FilterFn, IsConstantFn, IsConstantOpts, MinMaxFn, MinMaxResult, ScalarAtFn,
12    SliceFn, SumFn, TakeFn, ToArrowFn, UncompressedSizeFn, filter, is_constant_opts, min_max,
13    scalar_at, slice, sum, take, uncompressed_size,
14};
15use crate::variants::ExtensionArrayTrait;
16use crate::vtable::ComputeVTable;
17use crate::{Array, ArrayRef};
18
19impl ComputeVTable for ExtensionEncoding {
20    fn cast_fn(&self) -> Option<&dyn CastFn<&dyn Array>> {
21        // It's not possible to cast an extension array to another type.
22        // TODO(ngates): we should allow some extension arrays to implement a callback
23        //  to support this
24        None
25    }
26
27    fn compare_fn(&self) -> Option<&dyn CompareFn<&dyn Array>> {
28        Some(self)
29    }
30
31    fn filter_fn(&self) -> Option<&dyn FilterFn<&dyn Array>> {
32        Some(self)
33    }
34
35    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
36        Some(self)
37    }
38
39    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
40        Some(self)
41    }
42
43    fn slice_fn(&self) -> Option<&dyn SliceFn<&dyn Array>> {
44        Some(self)
45    }
46
47    fn sum_fn(&self) -> Option<&dyn SumFn<&dyn Array>> {
48        Some(self)
49    }
50
51    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
52        Some(self)
53    }
54
55    fn to_arrow_fn(&self) -> Option<&dyn ToArrowFn<&dyn Array>> {
56        Some(self)
57    }
58
59    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
60        Some(self)
61    }
62
63    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
64        Some(self)
65    }
66}
67
68impl FilterFn<&ExtensionArray> for ExtensionEncoding {
69    fn filter(&self, array: &ExtensionArray, mask: &Mask) -> VortexResult<ArrayRef> {
70        Ok(
71            ExtensionArray::new(array.ext_dtype().clone(), filter(array.storage(), mask)?)
72                .into_array(),
73        )
74    }
75}
76
77impl ScalarAtFn<&ExtensionArray> for ExtensionEncoding {
78    fn scalar_at(&self, array: &ExtensionArray, index: usize) -> VortexResult<Scalar> {
79        Ok(Scalar::extension(
80            array.ext_dtype().clone(),
81            scalar_at(array.storage(), index)?,
82        ))
83    }
84}
85
86impl SliceFn<&ExtensionArray> for ExtensionEncoding {
87    fn slice(&self, array: &ExtensionArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
88        Ok(ExtensionArray::new(
89            array.ext_dtype().clone(),
90            slice(array.storage(), start, stop)?,
91        )
92        .into_array())
93    }
94}
95
96impl SumFn<&ExtensionArray> for ExtensionEncoding {
97    fn sum(&self, array: &ExtensionArray) -> VortexResult<Scalar> {
98        sum(array.storage())
99    }
100}
101
102impl TakeFn<&ExtensionArray> for ExtensionEncoding {
103    fn take(&self, array: &ExtensionArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
104        Ok(
105            ExtensionArray::new(array.ext_dtype().clone(), take(array.storage(), indices)?)
106                .into_array(),
107        )
108    }
109}
110
111impl MinMaxFn<&ExtensionArray> for ExtensionEncoding {
112    fn min_max(&self, array: &ExtensionArray) -> VortexResult<Option<MinMaxResult>> {
113        Ok(
114            min_max(array.storage())?.map(|MinMaxResult { min, max }| MinMaxResult {
115                min: Scalar::extension(array.ext_dtype().clone(), min),
116                max: Scalar::extension(array.ext_dtype().clone(), max),
117            }),
118        )
119    }
120}
121
122impl IsConstantFn<&ExtensionArray> for ExtensionEncoding {
123    fn is_constant(
124        &self,
125        array: &ExtensionArray,
126        opts: &IsConstantOpts,
127    ) -> VortexResult<Option<bool>> {
128        is_constant_opts(array.storage(), opts).map(Some)
129    }
130}
131
132impl UncompressedSizeFn<&ExtensionArray> for ExtensionEncoding {
133    fn uncompressed_size(&self, array: &ExtensionArray) -> VortexResult<usize> {
134        uncompressed_size(array.storage())
135    }
136}
vortex_array/arrays/extension/compute/mod.rs

vortex_array/arrays/extension/compute/
mod.rs
