vortex_array/expr/exprs/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::Array;
11use crate::ArrayRef;
12use crate::ExecutionCtx;
13use crate::arrays::ExactScalarFn;
14use crate::arrays::ScalarFnArrayView;
15use crate::arrays::ScalarFnVTable;
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: &dyn Array,
27        upper: &dyn Array,
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: &dyn Array,
39        upper: &dyn Array,
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].as_ref();
70        let upper = children[2].as_ref();
71        if let Some(result) = precondition(&**array, lower, upper)? {
72            return Ok(Some(result));
73        }
74        <V as BetweenReduce>::between(array, lower, upper, parent.options)
75    }
76}
77
78/// Adapts a [`BetweenKernel`] impl into an [`ExecuteParentKernel`] for `ScalarFnArray(Between, ...)`.
79#[derive(Default, Debug)]
80pub struct BetweenExecuteAdaptor<V>(pub V);
81
82impl<V> ExecuteParentKernel<V> for BetweenExecuteAdaptor<V>
83where
84    V: BetweenKernel,
85{
86    type Parent = ExactScalarFn<Between>;
87
88    fn execute_parent(
89        &self,
90        array: &V::Array,
91        parent: ScalarFnArrayView<'_, Between>,
92        child_idx: usize,
93        ctx: &mut ExecutionCtx,
94    ) -> VortexResult<Option<ArrayRef>> {
95        // Only process the main array child (index 0), not lower (1) or upper (2).
96        if child_idx != 0 {
97            return Ok(None);
98        }
99        let scalar_fn_array = parent
100            .as_opt::<ScalarFnVTable>()
101            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
102        let children = scalar_fn_array.children();
103        let lower = children[1].as_ref();
104        let upper = children[2].as_ref();
105        if let Some(result) = precondition(&**array, lower, upper)? {
106            return Ok(Some(result));
107        }
108        <V as BetweenKernel>::between(array, lower, upper, parent.options, ctx)
109    }
110}
vortex_array/expr/exprs/between/kernel.rs

vortex_array/expr/exprs/between/
kernel.rs
