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    FilterKernel, FilterKernelAdapter, IsConstantFn, IsConstantOpts, IsSortedFn, MinMaxFn,
12    MinMaxResult, ScalarAtFn, SliceFn, SumKernel, SumKernelAdapter, 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, ArrayRef, register_kernel};
19
20impl ComputeVTable for ExtensionEncoding {
21    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
22        Some(self)
23    }
24
25    fn is_sorted_fn(&self) -> Option<&dyn IsSortedFn<&dyn Array>> {
26        Some(self)
27    }
28
29    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
30        Some(self)
31    }
32
33    fn slice_fn(&self) -> Option<&dyn SliceFn<&dyn Array>> {
34        Some(self)
35    }
36
37    fn take_fn(&self) -> Option<&dyn TakeFn<&dyn Array>> {
38        Some(self)
39    }
40
41    fn to_arrow_fn(&self) -> Option<&dyn ToArrowFn<&dyn Array>> {
42        Some(self)
43    }
44
45    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
46        Some(self)
47    }
48
49    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
50        Some(self)
51    }
52}
53
54impl FilterKernel for ExtensionEncoding {
55    fn filter(&self, array: &ExtensionArray, mask: &Mask) -> VortexResult<ArrayRef> {
56        Ok(
57            ExtensionArray::new(array.ext_dtype().clone(), filter(array.storage(), mask)?)
58                .into_array(),
59        )
60    }
61}
62
63register_kernel!(FilterKernelAdapter(ExtensionEncoding).lift());
64
65impl ScalarAtFn<&ExtensionArray> for ExtensionEncoding {
66    fn scalar_at(&self, array: &ExtensionArray, index: usize) -> VortexResult<Scalar> {
67        Ok(Scalar::extension(
68            array.ext_dtype().clone(),
69            scalar_at(array.storage(), index)?,
70        ))
71    }
72}
73
74impl SliceFn<&ExtensionArray> for ExtensionEncoding {
75    fn slice(&self, array: &ExtensionArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
76        Ok(ExtensionArray::new(
77            array.ext_dtype().clone(),
78            slice(array.storage(), start, stop)?,
79        )
80        .into_array())
81    }
82}
83
84impl SumKernel for ExtensionEncoding {
85    fn sum(&self, array: &ExtensionArray) -> VortexResult<Scalar> {
86        sum(array.storage())
87    }
88}
89
90register_kernel!(SumKernelAdapter(ExtensionEncoding).lift());
91
92impl TakeFn<&ExtensionArray> for ExtensionEncoding {
93    fn take(&self, array: &ExtensionArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
94        Ok(
95            ExtensionArray::new(array.ext_dtype().clone(), take(array.storage(), indices)?)
96                .into_array(),
97        )
98    }
99}
100
101impl MinMaxFn<&ExtensionArray> for ExtensionEncoding {
102    fn min_max(&self, array: &ExtensionArray) -> VortexResult<Option<MinMaxResult>> {
103        Ok(
104            min_max(array.storage())?.map(|MinMaxResult { min, max }| MinMaxResult {
105                min: Scalar::extension(array.ext_dtype().clone(), min),
106                max: Scalar::extension(array.ext_dtype().clone(), max),
107            }),
108        )
109    }
110}
111
112impl IsConstantFn<&ExtensionArray> for ExtensionEncoding {
113    fn is_constant(
114        &self,
115        array: &ExtensionArray,
116        opts: &IsConstantOpts,
117    ) -> VortexResult<Option<bool>> {
118        is_constant_opts(array.storage(), opts).map(Some)
119    }
120}
121
122impl UncompressedSizeFn<&ExtensionArray> for ExtensionEncoding {
123    fn uncompressed_size(&self, array: &ExtensionArray) -> VortexResult<usize> {
124        uncompressed_size(array.storage())
125    }
126}
127
128impl IsSortedFn<&ExtensionArray> for ExtensionEncoding {
129    fn is_sorted(&self, array: &ExtensionArray) -> VortexResult<bool> {
130        is_sorted(array.storage())
131    }
132
133    fn is_strict_sorted(&self, array: &ExtensionArray) -> VortexResult<bool> {
134        is_strict_sorted(array.storage())
135    }
136}
vortex_array/arrays/extension/compute/mod.rs

vortex_array/arrays/extension/compute/
mod.rs
