vortex_array/arrays/extension/compute/
mod.rs

1mod 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());