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

vortex_array/scalar_fn/fns/between/
kernel.rs
