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

vortex_array/expr/exprs/zip/
kernel.rs
