Add MaskFn (#5606)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/arrays/masked/vtable/mod.rs

@@ -7,11 +7,11 @@ mod operations;
 mod validity;
 
 use vortex_buffer::BufferHandle;
-use vortex_compute::mask::MaskValidity;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_vector::Vector;
+use vortex_vector::VectorOps;
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
@@ -106,8 +106,9 @@ impl VTable for MaskedVTable {
     }
 
     fn batch_execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
-        let vector = array.child().batch_execute(ctx)?;
-        Ok(MaskValidity::mask_validity(vector, &array.validity_mask()))
+        let mut vector = array.child().batch_execute(ctx)?;
+        vector.mask_validity(&array.validity_mask());
+        Ok(vector)
     }
 }
 

vortex-array/src/expr/exprs/get_item/mod.rs

@@ -7,7 +7,6 @@ use std::fmt::Formatter;
 use std::ops::Not;
 
 use prost::Message;
-use vortex_compute::mask::MaskValidity;
 use vortex_dtype::DType;
 use vortex_dtype::FieldName;
 use vortex_dtype::FieldPath;
@@ -145,8 +144,8 @@ impl VTable for GetItem {
             .into_struct();
 
         // We must intersect the validity with that of the parent struct
-        let field = struct_vector.fields()[field_idx].clone();
-        let field = MaskValidity::mask_validity(field, struct_vector.validity());
+        let mut field = struct_vector.fields()[field_idx].clone();
+        field.mask_validity(struct_vector.validity());
 
         Ok(field)
     }

vortex-array/src/expr/functions/vtable.rs

@@ -110,15 +110,15 @@ pub trait VTable: 'static + Send + Sync + Sized {
 /// The arity (number of arguments) of a function.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Arity {
-    Fixed(usize),
+    Exact(usize),
     Variadic { min: usize, max: Option<usize> },
 }
 
 impl Arity {
     /// Whether the given argument count matches this arity.
     pub fn matches(&self, arg_count: usize) -> bool {
         match self {
-            Arity::Fixed(m) => *m == arg_count,
+            Arity::Exact(m) => *m == arg_count,
             Arity::Variadic { min, max } => {
                 if arg_count < *min {
                     return false;

vortex-array/src/scalar_fns/cast/array.rs

@@ -41,7 +41,7 @@ mod test {
     use crate::arrays::ConstantArray;
     use crate::arrays::ConstantVTable;
     use crate::optimizer::ArrayOptimizer;
-    use crate::scalar_fns::BuiltinScalarFns;
+    use crate::scalar_fns::ArrayBuiltins;
 
     #[test]
     fn test_same_dtype() -> VortexResult<()> {

vortex-array/src/scalar_fns/cast/mod.rs

@@ -21,7 +21,7 @@ use crate::expr::functions::FunctionId;
 use crate::expr::functions::NullHandling;
 use crate::expr::functions::VTable;
 use crate::expr::stats::Stat;
-use crate::scalar_fns::BuiltinScalarFns;
+use crate::scalar_fns::ExprBuiltins;
 
 pub struct CastFn;
 impl VTable for CastFn {
@@ -49,7 +49,7 @@ impl VTable for CastFn {
     }
 
     fn arity(&self, _options: &DType) -> Arity {
-        Arity::Fixed(1)
+        Arity::Exact(1)
     }
 
     fn null_handling(&self, _options: &DType) -> NullHandling {

vortex-array/src/scalar_fns/is_null/mod.rs

@@ -30,7 +30,7 @@ impl VTable for IsNullFn {
     }
 
     fn arity(&self, _: &Self::Options) -> Arity {
-        Arity::Fixed(1)
+        Arity::Exact(1)
     }
 
     fn arg_name(&self, _: &Self::Options, arg_idx: usize) -> ArgName {

vortex-array/src/scalar_fns/mask/mod.rs

@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::DType;
+use vortex_dtype::Nullability;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_mask::Mask;
+use vortex_vector::BoolDatum;
+use vortex_vector::Datum;
+use vortex_vector::ScalarOps;
+use vortex_vector::VectorMutOps;
+use vortex_vector::VectorOps;
+
+use crate::expr::functions::ArgName;
+use crate::expr::functions::Arity;
+use crate::expr::functions::EmptyOptions;
+use crate::expr::functions::ExecutionArgs;
+use crate::expr::functions::FunctionId;
+use crate::expr::functions::VTable;
+
+/// A function that intersects the validity of an array using another array as a mask.
+///
+/// Where the `mask` array is true, the corresponding v
+pub struct MaskFn;
+impl VTable for MaskFn {
+    type Options = EmptyOptions;
+
+    fn id(&self) -> FunctionId {
+        FunctionId::from("vortex.mask")
+    }
+
+    fn arity(&self, _options: &Self::Options) -> Arity {
+        Arity::Exact(2)
+    }
+
+    fn arg_name(&self, _options: &Self::Options, arg_idx: usize) -> ArgName {
+        match arg_idx {
+            0 => ArgName::from("input"),
+            1 => ArgName::from("mask"),
+            _ => unreachable!("unknown"),
+        }
+    }
+
+    fn return_dtype(&self, _options: &Self::Options, arg_types: &[DType]) -> VortexResult<DType> {
+        vortex_ensure!(
+            arg_types[1] == DType::Bool(Nullability::NonNullable),
+            "The mask argument to 'mask' must be a non-nullable boolean array, got {}",
+            arg_types[1]
+        );
+        Ok(arg_types[0].as_nullable())
+    }
+
+    fn execute(&self, _options: &Self::Options, args: &ExecutionArgs) -> VortexResult<Datum> {
+        let input = args.input_datums(0).clone();
+        let mask = args.input_datums(1).clone().into_bool();
+        match (input, mask) {
+            (Datum::Scalar(input), BoolDatum::Scalar(mask)) => {
+                let mut result = input;
+                result.mask_validity(mask.value().vortex_expect("mask is non-nullable"));
+                Ok(Datum::Scalar(result))
+            }
+            (Datum::Scalar(input), BoolDatum::Vector(mask)) => {
+                let mut result = input.repeat(args.row_count()).freeze();
+                result.mask_validity(&Mask::from(mask.into_bits()));
+                Ok(Datum::Vector(result))
+            }
+            (Datum::Vector(input_array), BoolDatum::Scalar(mask)) => {
+                let mut result = input_array;
+                result.mask_validity(&Mask::new(
+                    args.row_count(),
+                    mask.value().vortex_expect("mask is non-nullable"),
+                ));
+                Ok(Datum::Vector(result))
+            }
+            (Datum::Vector(input_array), BoolDatum::Vector(mask)) => {
+                let mut result = input_array;
+                result.mask_validity(&Mask::from(mask.into_bits()));
+                Ok(Datum::Vector(result))
+            }
+        }
+    }
+}

vortex-array/src/scalar_fns/mod.rs

@@ -17,25 +17,78 @@ use crate::ArrayRef;
 use crate::arrays::ScalarFnArrayExt;
 use crate::expr::Expression;
 use crate::expr::ScalarFnExprExt;
+use crate::expr::functions::EmptyOptions;
 
 pub mod cast;
 pub mod is_null;
+pub mod mask;
 pub mod not;
 
 /// A collection of built-in scalar functions that can be applied to expressions or arrays.
-pub trait BuiltinScalarFns: Sized {
+pub trait ExprBuiltins: Sized {
     /// Cast to the given data type.
-    fn cast(&self, dtype: DType) -> VortexResult<Self>;
+    fn cast(&self, dtype: DType) -> VortexResult<Expression>;
+
+    /// Is null check.
+    fn is_null(&self) -> VortexResult<Expression>;
+
+    /// Boolean negation.
+    fn not(&self) -> VortexResult<Expression>;
+
+    /// Mask the expression using the given boolean mask.
+    /// The resulting expression's validity is the intersection of the original expression's
+    /// validity.
+    fn mask(&self, mask: Expression) -> VortexResult<Expression>;
 }
 
-impl BuiltinScalarFns for Expression {
+impl ExprBuiltins for Expression {
     fn cast(&self, dtype: DType) -> VortexResult<Expression> {
         cast::CastFn.try_new_expr(dtype, [self.clone()])
     }
+
+    fn is_null(&self) -> VortexResult<Expression> {
+        is_null::IsNullFn.try_new_expr(EmptyOptions, [self.clone()])
+    }
+
+    fn not(&self) -> VortexResult<Expression> {
+        not::NotFn.try_new_expr(EmptyOptions, [self.clone()])
+    }
+
+    fn mask(&self, mask: Expression) -> VortexResult<Expression> {
+        mask::MaskFn.try_new_expr(EmptyOptions, [self.clone(), mask])
+    }
 }
 
-impl BuiltinScalarFns for ArrayRef {
-    fn cast(&self, dtype: DType) -> VortexResult<Self> {
+pub trait ArrayBuiltins: Sized {
+    /// Cast to the given data type.
+    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef>;
+
+    /// Is null check.
+    fn is_null(&self) -> VortexResult<ArrayRef>;
+
+    /// Boolean negation.
+    fn not(&self) -> VortexResult<ArrayRef>;
+
+    /// Mask the array using the given boolean mask.
+    /// The resulting array's validity is the intersection of the original array's validity
+    /// and the mask's validity.
+    fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef>;
+}
+
+impl ArrayBuiltins for ArrayRef {
+    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef> {
         cast::CastFn.try_new_array(self.len(), dtype, [self.clone()])
     }
+
+    fn is_null(&self) -> VortexResult<ArrayRef> {
+        is_null::IsNullFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
+    }
+
+    fn not(&self) -> VortexResult<ArrayRef> {
+        not::NotFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
+    }
+
+    fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef> {
+        mask::MaskFn.try_new_array(self.len(), EmptyOptions, [self.clone(), mask.clone()])
+    }
 }

vortex-array/src/scalar_fns/not/mod.rs

@@ -28,7 +28,7 @@ impl VTable for NotFn {
     }
 
     fn arity(&self, _: &Self::Options) -> Arity {
-        Arity::Fixed(1)
+        Arity::Exact(1)
     }
 
     fn null_handling(&self, _options: &Self::Options) -> NullHandling {

vortex-compute/src/lib.rs

@@ -15,5 +15,4 @@ pub mod comparison;
 pub mod expand;
 pub mod filter;
 pub mod logical;
-pub mod mask;
 pub mod take;