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

vortex_array/scalar_fn/fns/mask/
mod.rs
