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

vortex_array/scalar_fn/fns/mask/
mod.rs
