vortex_array/scalar_fn/fns/mask/
kernel.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Encoding-specific mask reduction and execution adapters.
5
6use vortex_error::VortexResult;
7use vortex_error::vortex_err;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::array::ArrayView;
12use crate::array::VTable;
13use crate::arrays::Bool;
14use crate::arrays::Constant;
15use crate::arrays::scalar_fn::ExactScalarFn;
16use crate::arrays::scalar_fn::ScalarFnArrayView;
17use crate::kernel::ExecuteParentKernel;
18use crate::optimizer::rules::ArrayParentReduceRule;
19use crate::scalar_fn::fns::mask::Mask as MaskExpr;
20
21/// Mask an array without reading buffers.
22///
23/// This trait is for mask implementations that can operate purely on array metadata and
24/// structure without needing to read or execute on the underlying buffers. Implementations
25/// should return `None` if masking requires buffer access.
26///
27/// The `mask` parameter is a boolean array where true=keep/valid, false=null-out.
28///
29/// # Preconditions
30///
31/// The mask has non-nullable Boolean dtype and the same length as the array. It may be lazy.
32pub trait MaskReduce: VTable {
33    /// Whether reading this encoding's validity requires only metadata access.
34    ///
35    /// Enables lazy masks on all-valid inputs. Only opt in when `array.validity()` cannot read
36    /// buffers or execute children, including for nullable inputs.
37    const VALIDITY_IS_METADATA_ONLY: bool = false;
38
39    fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
40}
41
42/// Mask an array, potentially reading buffers.
43///
44/// Unlike [`MaskReduce`], this trait is for mask implementations that may need to read
45/// and execute on the underlying buffers to produce the masked result.
46///
47/// The `mask` parameter is a boolean array where true=keep/valid, false=null-out.
48///
49/// # Preconditions
50///
51/// The mask is guaranteed to have the same length as the array. Trivial cases
52/// (`AllValid`, `AllInvalid`, `NonNullable`) are handled by the caller before dispatch.
53pub trait MaskKernel: VTable {
54    fn mask(
55        array: ArrayView<'_, Self>,
56        mask: &ArrayRef,
57        ctx: &mut ExecutionCtx,
58    ) -> VortexResult<Option<ArrayRef>>;
59}
60
61/// Adaptor that wraps a [`MaskReduce`] impl as an [`ArrayParentReduceRule`].
62#[derive(Default, Debug)]
63pub struct MaskReduceAdaptor<V>(pub V);
64
65impl<V> ArrayParentReduceRule<V> for MaskReduceAdaptor<V>
66where
67    V: MaskReduce,
68{
69    type Parent = ExactScalarFn<MaskExpr>;
70
71    fn reduce_parent(
72        &self,
73        array: ArrayView<'_, V>,
74        parent: ScalarFnArrayView<'_, MaskExpr>,
75        child_idx: usize,
76    ) -> VortexResult<Option<ArrayRef>> {
77        // Only reduce the input child (index 0), not the mask child (index 1).
78        if child_idx != 0 {
79            return Ok(None);
80        }
81
82        let parent_ref: ArrayRef = (*parent).clone();
83        let mask_child = parent_ref
84            .nth_child(1)
85            .ok_or_else(|| vortex_err!("Mask expression must have 2 children"))?;
86
87        if mask_child.as_opt::<Bool>().is_none() && mask_child.as_opt::<Constant>().is_none() {
88            let can_attach_mask =
89                V::VALIDITY_IS_METADATA_ONLY && array.validity()?.definitely_no_nulls();
90
91            if !can_attach_mask {
92                return Ok(None);
93            }
94        }
95
96        <V as MaskReduce>::mask(array, mask_child)
97    }
98}
99
100/// Adaptor that wraps a [`MaskKernel`] impl as an [`ExecuteParentKernel`].
101#[derive(Default, Debug)]
102pub struct MaskExecuteAdaptor<V>(pub V);
103
104impl<V> ExecuteParentKernel<V> for MaskExecuteAdaptor<V>
105where
106    V: MaskKernel,
107{
108    type Parent = ExactScalarFn<MaskExpr>;
109
110    fn execute_parent(
111        &self,
112        array: ArrayView<'_, V>,
113        parent: ScalarFnArrayView<'_, MaskExpr>,
114        child_idx: usize,
115        ctx: &mut ExecutionCtx,
116    ) -> VortexResult<Option<ArrayRef>> {
117        // Only execute the input child (index 0), not the mask child (index 1).
118        if child_idx != 0 {
119            return Ok(None);
120        }
121        let mask_child = parent
122            .nth_child(1)
123            .ok_or_else(|| vortex_err!("Mask expression must have 2 children"))?;
124        <V as MaskKernel>::mask(array, mask_child, ctx)
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use rstest::rstest;
131    use vortex_buffer::buffer;
132    use vortex_error::VortexResult;
133
134    use crate::ArrayEq;
135    use crate::ArrayRef;
136    use crate::EqMode;
137    use crate::IntoArray;
138    use crate::arrays::BoolArray;
139    use crate::arrays::ConstantArray;
140    use crate::arrays::Primitive;
141    use crate::arrays::PrimitiveArray;
142    use crate::arrays::ScalarFn;
143    use crate::arrays::scalar_fn::ExactScalarFn;
144    use crate::assert_arrays_eq;
145    use crate::builtins::ArrayBuiltins;
146    use crate::dtype::Nullability;
147    use crate::executor::VortexSessionExecute;
148    use crate::optimizer::ArrayOptimizer;
149    use crate::scalar::Scalar;
150    use crate::scalar_fn::fns::binary::Binary;
151    use crate::scalar_fn::fns::mask::Mask as MaskExpr;
152    use crate::scalar_fn::fns::operators::Operator;
153    use crate::validity::Validity;
154
155    fn lazy_mask() -> VortexResult<ArrayRef> {
156        Binary::try_new(
157            BoolArray::from_iter([true, true, false]).into_array(),
158            BoolArray::from_iter([true, false, true]).into_array(),
159            Operator::And,
160        )
161        .map(IntoArray::into_array)
162    }
163
164    #[rstest]
165    #[case::non_nullable(Validity::NonNullable)]
166    #[case::all_valid(Validity::AllValid)]
167    fn all_valid_input_attaches_lazy_mask(#[case] validity: Validity) -> VortexResult<()> {
168        let input = PrimitiveArray::new(buffer![1i32, 2, 3], validity);
169        let mask = lazy_mask()?;
170        let output = input.clone().into_array().mask(mask.clone())?;
171        let primitive = output.as_::<Primitive>();
172
173        assert!(
174            primitive
175                .buffer_handle()
176                .array_eq(input.buffer_handle(), EqMode::Ptr)
177        );
178        assert_eq!(output.dtype(), &input.dtype().as_nullable());
179        assert_eq!(output.len(), input.len());
180        let Validity::Array(validity) = output.validity()? else {
181            panic!("mask must remain array-backed");
182        };
183        assert!(ArrayRef::ptr_eq(&validity, &mask));
184
185        Ok(())
186    }
187
188    #[rstest]
189    #[case::partially_valid(false)]
190    #[case::all_true_bitmap(true)]
191    fn array_backed_validity_keeps_lazy_mask(#[case] all_true: bool) -> VortexResult<()> {
192        let validity = BoolArray::from_iter([true, all_true, true]).into_array();
193        let input =
194            PrimitiveArray::new(buffer![1i32, 2, 3], Validity::Array(validity)).into_array();
195        let output = input.mask(lazy_mask()?)?;
196        assert!(output.is::<ExactScalarFn<MaskExpr>>());
197
198        Ok(())
199    }
200
201    /// A constant Boolean mask child must take the metadata-only reduction path (pushing into the
202    /// input encoding) rather than surviving as a `ScalarFn` wrapper that falls through to
203    /// execution. Asserting the optimized encoding makes this fail before the adaptor accepts
204    /// `Constant` masks, not just verifying values that could pass through the execution fallback.
205    #[rstest]
206    #[case(true)]
207    #[case(false)]
208    fn constant_mask_reduces_into_input(#[case] mask_value: bool) -> VortexResult<()> {
209        let input = buffer![1i32, 2, 3, 4, 5].into_array();
210        let mask = ConstantArray::new(
211            Scalar::bool(mask_value, Nullability::NonNullable),
212            input.len(),
213        )
214        .into_array();
215
216        let masked = MaskExpr::try_new(input, mask)?.into_array();
217        assert!(
218            masked.is::<ScalarFn>(),
219            "expected an un-optimized ScalarFn wrapper before optimization"
220        );
221
222        let optimized = masked.optimize()?;
223        assert!(
224            !optimized.is::<ScalarFn>(),
225            "constant mask should not fall through to execution, got {}",
226            optimized.encoding_id()
227        );
228        assert!(
229            optimized.is::<Primitive>(),
230            "constant mask should reduce into the Primitive input, got {}",
231            optimized.encoding_id()
232        );
233
234        let mut ctx = crate::array_session().create_execution_ctx();
235        let expected = if mask_value {
236            PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4), Some(5)])
237        } else {
238            PrimitiveArray::from_option_iter([None::<i32>, None, None, None, None])
239        };
240        assert_arrays_eq!(optimized, expected, &mut ctx);
241
242        Ok(())
243    }
244}
vortex_array/scalar_fn/fns/mask/kernel.rs

vortex_array/scalar_fn/fns/mask/
kernel.rs
