vortex_array/arrays/extension/compute/
mod.rs

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

vortex_array/arrays/extension/compute/
mod.rs
