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

vortex_array/scalar_fn/fns/between/
kernel.rs
