vortex_array/arrays/extension/compute/
mod.rs1mod compare;
2
3use vortex_error::VortexResult;
4use vortex_mask::Mask;
5use vortex_scalar::Scalar;
6
7use crate::arrays::ExtensionVTable;
8use crate::arrays::extension::ExtensionArray;
9use crate::compute::{
10 FilterKernel, FilterKernelAdapter, IsConstantKernel, IsConstantKernelAdapter, IsConstantOpts,
11 IsSortedKernel, IsSortedKernelAdapter, MinMaxKernel, MinMaxKernelAdapter, MinMaxResult,
12 SumKernel, SumKernelAdapter, TakeKernel, TakeKernelAdapter, filter, is_constant_opts,
13 is_sorted, is_strict_sorted, min_max, sum, take,
14};
15use crate::{Array, ArrayRef, IntoArray, register_kernel};
16
17impl FilterKernel for ExtensionVTable {
18 fn filter(&self, array: &ExtensionArray, mask: &Mask) -> VortexResult<ArrayRef> {
19 Ok(
20 ExtensionArray::new(array.ext_dtype().clone(), filter(array.storage(), mask)?)
21 .into_array(),
22 )
23 }
24}
25
26register_kernel!(FilterKernelAdapter(ExtensionVTable).lift());
27
28impl SumKernel for ExtensionVTable {
29 fn sum(&self, array: &ExtensionArray) -> VortexResult<Scalar> {
30 sum(array.storage())
31 }
32}
33
34register_kernel!(SumKernelAdapter(ExtensionVTable).lift());
35
36impl TakeKernel for ExtensionVTable {
37 fn take(&self, array: &ExtensionArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
38 Ok(
39 ExtensionArray::new(array.ext_dtype().clone(), take(array.storage(), indices)?)
40 .into_array(),
41 )
42 }
43}
44
45register_kernel!(TakeKernelAdapter(ExtensionVTable).lift());
46
47impl MinMaxKernel for ExtensionVTable {
48 fn min_max(&self, array: &ExtensionArray) -> VortexResult<Option<MinMaxResult>> {
49 Ok(
50 min_max(array.storage())?.map(|MinMaxResult { min, max }| MinMaxResult {
51 min: Scalar::extension(array.ext_dtype().clone(), min),
52 max: Scalar::extension(array.ext_dtype().clone(), max),
53 }),
54 )
55 }
56}
57
58register_kernel!(MinMaxKernelAdapter(ExtensionVTable).lift());
59
60impl IsConstantKernel for ExtensionVTable {
61 fn is_constant(
62 &self,
63 array: &ExtensionArray,
64 opts: &IsConstantOpts,
65 ) -> VortexResult<Option<bool>> {
66 is_constant_opts(array.storage(), opts)
67 }
68}
69
70register_kernel!(IsConstantKernelAdapter(ExtensionVTable).lift());
71
72impl IsSortedKernel for ExtensionVTable {
73 fn is_sorted(&self, array: &ExtensionArray) -> VortexResult<bool> {
74 is_sorted(array.storage())
75 }
76
77 fn is_strict_sorted(&self, array: &ExtensionArray) -> VortexResult<bool> {
78 is_strict_sorted(array.storage())
79 }
80}
81
82register_kernel!(IsSortedKernelAdapter(ExtensionVTable).lift());