vortex_array/scalar_fn/fns/zip/
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 crate::ArrayRef;
8use crate::ExecutionCtx;
9use crate::arrays::ExactScalarFn;
10use crate::arrays::ScalarFnArrayView;
11use crate::arrays::ScalarFnVTable;
12use crate::kernel::ExecuteParentKernel;
13use crate::optimizer::rules::ArrayParentReduceRule;
14use crate::scalar_fn::fns::zip::Zip as ZipExpr;
15use crate::vtable::VTable;
16
17/// Zip two arrays using a mask without reading buffers.
18///
19/// This trait is for zip implementations that can operate purely on array metadata
20/// and structure without needing to read or execute on the underlying buffers.
21/// Implementations should return `None` if the operation requires buffer access.
22///
23/// Dispatch is on child 0 (if_true). The `if_false` and `mask` are extracted from
24/// the parent `ScalarFnArray`.
25pub trait ZipReduce: VTable {
26    fn zip(
27        array: &Self::Array,
28        if_false: &ArrayRef,
29        mask: &ArrayRef,
30    ) -> VortexResult<Option<ArrayRef>>;
31}
32
33/// Zip two arrays using a mask, potentially reading buffers.
34///
35/// Unlike [`ZipReduce`], this trait is for zip implementations that may need
36/// to read and execute on the underlying buffers to produce the result.
37///
38/// Dispatch is on child 0 (if_true). The `if_false` and `mask` are extracted from
39/// the parent `ScalarFnArray`.
40pub trait ZipKernel: VTable {
41    fn zip(
42        array: &Self::Array,
43        if_false: &ArrayRef,
44        mask: &ArrayRef,
45        ctx: &mut ExecutionCtx,
46    ) -> VortexResult<Option<ArrayRef>>;
47}
48
49/// Adaptor that wraps a [`ZipReduce`] impl as an [`ArrayParentReduceRule`].
50#[derive(Default, Debug)]
51pub struct ZipReduceAdaptor<V>(pub V);
52
53impl<V> ArrayParentReduceRule<V> for ZipReduceAdaptor<V>
54where
55    V: ZipReduce,
56{
57    type Parent = ExactScalarFn<ZipExpr>;
58
59    fn reduce_parent(
60        &self,
61        array: &V::Array,
62        parent: ScalarFnArrayView<'_, ZipExpr>,
63        child_idx: usize,
64    ) -> VortexResult<Option<ArrayRef>> {
65        if child_idx != 0 {
66            return Ok(None);
67        }
68        let scalar_fn_array = parent
69            .as_opt::<ScalarFnVTable>()
70            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
71        let children = scalar_fn_array.children();
72        let if_false = &children[1];
73        let mask_array = &children[2];
74        <V as ZipReduce>::zip(array, if_false, mask_array)
75    }
76}
77
78/// Adaptor that wraps a [`ZipKernel`] impl as an [`ExecuteParentKernel`].
79#[derive(Default, Debug)]
80pub struct ZipExecuteAdaptor<V>(pub V);
81
82impl<V> ExecuteParentKernel<V> for ZipExecuteAdaptor<V>
83where
84    V: ZipKernel,
85{
86    type Parent = ExactScalarFn<ZipExpr>;
87
88    fn execute_parent(
89        &self,
90        array: &V::Array,
91        parent: ScalarFnArrayView<'_, ZipExpr>,
92        child_idx: usize,
93        ctx: &mut ExecutionCtx,
94    ) -> VortexResult<Option<ArrayRef>> {
95        if child_idx != 0 {
96            return Ok(None);
97        }
98        let scalar_fn_array = parent
99            .as_opt::<ScalarFnVTable>()
100            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
101        let children = scalar_fn_array.children();
102        let if_false = &children[1];
103        let mask_array = &children[2];
104        <V as ZipKernel>::zip(array, if_false, mask_array, ctx)
105    }
106}
vortex_array/scalar_fn/fns/zip/kernel.rs

vortex_array/scalar_fn/fns/zip/
kernel.rs
