vortex_array/scalar_fn/fns/between/
kernel.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6
7use super::Between;
8use super::BetweenOptions;
9use super::precondition;
10use crate::ArrayRef;
11use crate::ExecutionCtx;
12use crate::array::ArrayView;
13use crate::array::VTable;
14use crate::arrays::ScalarFnVTable;
15use crate::arrays::scalar_fn::ExactScalarFn;
16use crate::arrays::scalar_fn::ScalarFnArrayExt;
17use crate::arrays::scalar_fn::ScalarFnArrayView;
18use crate::kernel::ExecuteParentKernel;
19use crate::optimizer::rules::ArrayParentReduceRule;
20
21/// Reduce rule for between: restructure the array without reading buffers.
22///
23/// Returns `Ok(None)` if the rule doesn't apply or buffer access is needed.
24pub trait BetweenReduce: VTable {
25    fn between(
26        array: ArrayView<'_, Self>,
27        lower: &ArrayRef,
28        upper: &ArrayRef,
29        options: &BetweenOptions,
30    ) -> VortexResult<Option<ArrayRef>>;
31}
32
33/// Execute kernel for between: perform the actual between check, potentially reading buffers.
34///
35/// Returns `Ok(None)` if this kernel cannot handle the given inputs.
36pub trait BetweenKernel: VTable {
37    fn between(
38        array: ArrayView<'_, Self>,
39        lower: &ArrayRef,
40        upper: &ArrayRef,
41        options: &BetweenOptions,
42        ctx: &mut ExecutionCtx,
43    ) -> VortexResult<Option<ArrayRef>>;
44}
45
46/// Adapts a [`BetweenReduce`] impl into an [`ArrayParentReduceRule`] for `ScalarFnArray(Between, ...)`.
47#[derive(Default, Debug)]
48pub struct BetweenReduceAdaptor<V>(pub V);
49
50impl<V> ArrayParentReduceRule<V> for BetweenReduceAdaptor<V>
51where
52    V: BetweenReduce,
53{
54    type Parent = ExactScalarFn<Between>;
55
56    fn reduce_parent(
57        &self,
58        array: ArrayView<'_, V>,
59        parent: ScalarFnArrayView<'_, Between>,
60        child_idx: usize,
61    ) -> VortexResult<Option<ArrayRef>> {
62        // Only process the main array child (index 0), not lower (1) or upper (2).
63        if child_idx != 0 {
64            return Ok(None);
65        }
66        let scalar_fn_array = parent
67            .as_opt::<ScalarFnVTable>()
68            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
69        let children = scalar_fn_array.children();
70        let lower = &children[1];
71        let upper = &children[2];
72        let arr = array.array().clone();
73        if let Some(result) = precondition(&arr, lower, upper)? {
74            return Ok(Some(result));
75        }
76        <V as BetweenReduce>::between(array, lower, upper, parent.options)
77    }
78}
79
80/// Adapts a [`BetweenKernel`] impl into an [`ExecuteParentKernel`] for `ScalarFnArray(Between, ...)`.
81#[derive(Default, Debug)]
82pub struct BetweenExecuteAdaptor<V>(pub V);
83
84impl<V> ExecuteParentKernel<V> for BetweenExecuteAdaptor<V>
85where
86    V: BetweenKernel,
87{
88    type Parent = ExactScalarFn<Between>;
89
90    fn execute_parent(
91        &self,
92        array: ArrayView<'_, V>,
93        parent: ScalarFnArrayView<'_, Between>,
94        child_idx: usize,
95        ctx: &mut ExecutionCtx,
96    ) -> VortexResult<Option<ArrayRef>> {
97        // Only process the main array child (index 0), not lower (1) or upper (2).
98        if child_idx != 0 {
99            return Ok(None);
100        }
101        let scalar_fn_array = parent
102            .as_opt::<ScalarFnVTable>()
103            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
104        let children = scalar_fn_array.children();
105        let lower = &children[1];
106        let upper = &children[2];
107        let arr = array.array().clone();
108        if let Some(result) = precondition(&arr, lower, upper)? {
109            return Ok(Some(result));
110        }
111        <V as BetweenKernel>::between(array, lower, upper, parent.options, ctx)
112    }
113}
vortex_array/scalar_fn/fns/between/kernel.rs

vortex_array/scalar_fn/fns/between/
kernel.rs
