vortex_array/scalar_fn/fns/mask/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod kernel;
5
6pub use kernel::*;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_ensure;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayRef;
14use crate::Canonical;
15use crate::ExecutionCtx;
16use crate::IntoArray;
17use crate::arrays::Constant;
18use crate::arrays::ConstantArray;
19use crate::arrays::ScalarFnArray;
20use crate::arrays::masked::mask_validity_canonical;
21use crate::builtins::ArrayBuiltins;
22use crate::child_to_validity;
23use crate::dtype::DType;
24use crate::dtype::Nullability;
25use crate::expr::Expression;
26use crate::expr::and;
27use crate::expr::lit;
28use crate::scalar::Scalar;
29use crate::scalar_fn::Arity;
30use crate::scalar_fn::ChildName;
31use crate::scalar_fn::EmptyOptions;
32use crate::scalar_fn::ExecutionArgs;
33use crate::scalar_fn::ScalarFnId;
34use crate::scalar_fn::ScalarFnVTable;
35use crate::scalar_fn::ScalarFnVTableExt;
36use crate::scalar_fn::SimplifyCtx;
37use crate::scalar_fn::fns::literal::Literal;
38
39/// An expression that masks an input based on a boolean mask.
40///
41/// Where the mask is true, the input value is retained; where the mask is false, the output is
42/// null. In other words, this performs an intersection of the input's validity with the mask.
43#[derive(Clone)]
44pub struct Mask;
45
46impl Mask {
47    /// Creates a lazy mask operation over `input`.
48    ///
49    /// # Errors
50    ///
51    /// Returns an error if the children have different lengths or `mask` is not non-nullable
52    /// boolean data.
53    pub fn try_new(input: ArrayRef, mask: ArrayRef) -> VortexResult<ScalarFnArray> {
54        ScalarFnArray::try_new(Mask.bind(EmptyOptions), vec![input, mask])
55    }
56}
57
58impl ScalarFnVTable for Mask {
59    type Options = EmptyOptions;
60
61    fn id(&self) -> ScalarFnId {
62        static ID: CachedId = CachedId::new("vortex.mask");
63        *ID
64    }
65
66    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
67        Ok(Some(vec![]))
68    }
69
70    fn deserialize(
71        &self,
72        _metadata: &[u8],
73        _session: &VortexSession,
74    ) -> VortexResult<Self::Options> {
75        Ok(EmptyOptions)
76    }
77
78    fn arity(&self, _options: &Self::Options) -> Arity {
79        Arity::Exact(2)
80    }
81
82    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
83        match child_idx {
84            0 => ChildName::from("input"),
85            1 => ChildName::from("mask"),
86            _ => unreachable!("Invalid child index {} for Mask expression", child_idx),
87        }
88    }
89
90    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
91        vortex_ensure!(
92            arg_dtypes[1] == DType::Bool(Nullability::NonNullable),
93            "The mask argument to 'mask' must be a non-nullable boolean array, got {}",
94            arg_dtypes[1]
95        );
96        Ok(arg_dtypes[0].as_nullable())
97    }
98
99    fn execute(
100        &self,
101        _options: &Self::Options,
102        args: &dyn ExecutionArgs,
103        ctx: &mut ExecutionCtx,
104    ) -> VortexResult<ArrayRef> {
105        let input = args.get(0)?;
106        let mask_array = args.get(1)?;
107
108        if let Some(result) = execute_constant(&input, &mask_array)? {
109            return Ok(result);
110        }
111
112        execute_canonical(input, mask_array, ctx)
113    }
114
115    fn simplify(
116        &self,
117        _options: &Self::Options,
118        expr: &Expression,
119        ctx: &dyn SimplifyCtx,
120    ) -> VortexResult<Option<Expression>> {
121        let Some(mask_lit) = expr.child(1).as_opt::<Literal>() else {
122            return Ok(None);
123        };
124
125        let mask_lit = mask_lit
126            .as_bool()
127            .value()
128            .vortex_expect("Mask must be non-nullable");
129
130        if mask_lit {
131            // Mask is all true, so the output is just the input.
132            Ok(Some(expr.child(0).clone()))
133        } else {
134            // Mask is all false, so the output is all nulls.
135            let input_dtype = ctx.return_dtype(expr.child(0))?;
136            Ok(Some(lit(Scalar::null(input_dtype.as_nullable()))))
137        }
138    }
139
140    fn validity(
141        &self,
142        _options: &Self::Options,
143        expression: &Expression,
144    ) -> VortexResult<Option<Expression>> {
145        Ok(Some(and(
146            expression.child(0).validity()?,
147            expression.child(1).clone(),
148        )))
149    }
150
151    fn is_strict(&self, _options: &Self::Options) -> bool {
152        true
153    }
154}
155
156/// Try to handle masking when at least one of the input or mask is a constant array.
157///
158/// Returns `Ok(Some(result))` if the constant case was handled, `Ok(None)` if not.
159fn execute_constant(input: &ArrayRef, mask_array: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
160    let len = input.len();
161
162    if let Some(constant_mask) = mask_array.as_opt::<Constant>() {
163        let mask_value = constant_mask.scalar().as_bool().value().unwrap_or(false);
164        return if mask_value {
165            input.cast(input.dtype().as_nullable()).map(Some)
166        } else {
167            Ok(Some(
168                ConstantArray::new(Scalar::null(input.dtype().as_nullable()), len).into_array(),
169            ))
170        };
171    }
172
173    if let Some(constant_input) = input.as_opt::<Constant>()
174        && constant_input.scalar().is_null()
175    {
176        return Ok(Some(
177            ConstantArray::new(Scalar::null(input.dtype().as_nullable()), len).into_array(),
178        ));
179    }
180
181    Ok(None)
182}
183
184/// Execute the mask by materializing both inputs to their canonical forms.
185fn execute_canonical(
186    input: ArrayRef,
187    mask_array: ArrayRef,
188    ctx: &mut ExecutionCtx,
189) -> VortexResult<ArrayRef> {
190    let validity = child_to_validity(Some(&mask_array), Nullability::Nullable);
191    let canonical = input.execute::<Canonical>(ctx)?;
192    Ok(mask_validity_canonical(canonical, validity, ctx)?.into_array())
193}
194
195#[cfg(test)]
196mod test {
197    use vortex_error::VortexExpect;
198
199    use crate::dtype::DType;
200    use crate::dtype::Nullability::Nullable;
201    use crate::dtype::PType;
202    use crate::expr::lit;
203    use crate::expr::mask;
204    use crate::scalar::Scalar;
205
206    #[test]
207    fn test_simplify() {
208        let input_expr = lit(42u32);
209        let true_mask_expr = lit(true);
210        let false_mask_expr = lit(false);
211
212        let mask_true_expr = mask(input_expr.clone(), true_mask_expr);
213        let simplified_true = mask_true_expr
214            .optimize(&DType::Null)
215            .vortex_expect("Simplification");
216        assert_eq!(&simplified_true, &input_expr);
217
218        let mask_false_expr = mask(input_expr, false_mask_expr);
219        let simplified_false = mask_false_expr
220            .optimize(&DType::Null)
221            .vortex_expect("Simplification");
222        let expected_null_expr = lit(Scalar::null(DType::Primitive(PType::U32, Nullable)));
223        assert_eq!(&simplified_false, &expected_null_expr);
224    }
225}
vortex_array/scalar_fn/fns/mask/mod.rs

vortex_array/scalar_fn/fns/mask/
mod.rs
