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

vortex_array/arrays/extension/compute/
mod.rs
