MaskFn

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/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 {

vortex-array/src/scalar_fns/mask.rs vortex-array/src/scalar_fns/mask/mod.rsvortex-array/src/scalar_fns/mask.rs renamed to vortex-array/src/scalar_fns/mask/mod.rs

@@ -60,7 +60,7 @@ impl VTable for MaskFn {
         let mask = args.input_datums(1).clone().into_bool();
         match (input, mask) {
             (Datum::Scalar(input), BoolDatum::Scalar(mask)) => {
-                let mut result = input.clone();
+                let mut result = input;
                 result.mask_validity(mask.value().vortex_expect("mask is non-nullable"));
                 Ok(Datum::Scalar(result))
             }
@@ -70,15 +70,15 @@ impl VTable for MaskFn {
                 Ok(Datum::Vector(result))
             }
             (Datum::Vector(input_array), BoolDatum::Scalar(mask)) => {
-                let mut result = input_array.clone();
+                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.clone();
+                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,26 +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;
-mod mask;
+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()])
+    }
 }