feat[scalar_fn]: is_null (#5601)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/expr/exprs/binary.rs

@@ -319,7 +319,7 @@ pub fn eq(lhs: Expression, rhs: Expression) -> Expression {
 /// ## Example usage
 ///
 /// ```
-/// # use vortex_array::arrays::{BoolArray, PrimitiveArray };
+/// # use vortex_array::arrays::{BoolArray, PrimitiveArray};
 /// # use vortex_array::{IntoArray, ToCanonical};
 /// # use vortex_array::validity::Validity;
 /// # use vortex_buffer::buffer;

vortex-array/src/expr/exprs/is_null.rs

@@ -4,6 +4,7 @@
 use std::fmt::Formatter;
 use std::ops::Not;
 
+use is_null::IsNullFn;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability;
 use vortex_error::VortexExpect;
@@ -24,12 +25,15 @@ use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::Expression;
 use crate::expr::ExpressionView;
+use crate::expr::ScalarFnExprExt;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
 use crate::expr::exprs::binary::eq;
 use crate::expr::exprs::literal::lit;
+use crate::expr::functions::EmptyOptions;
 use crate::expr::stats::Stat;
+use crate::scalar_fns::is_null;
 
 /// Expression that checks for null values.
 pub struct IsNull;
@@ -110,6 +114,10 @@ impl VTable for IsNull {
     fn is_fallible(&self, _instance: &Self::Instance) -> bool {
         false
     }
+
+    fn expr_v2(&self, view: &ExpressionView<Self>) -> VortexResult<Expression> {
+        ScalarFnExprExt::try_new_expr(&IsNullFn, EmptyOptions, view.children().clone())
+    }
 }
 
 /// Creates an expression that checks for null values.

vortex-array/src/expr/exprs/not.rs

@@ -17,8 +17,11 @@ use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::Expression;
 use crate::expr::ExpressionView;
+use crate::expr::ScalarFnExprExt;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
+use crate::expr::functions::EmptyOptions;
+use crate::scalar_fns::not;
 
 /// Expression that logically inverts boolean values.
 pub struct Not;
@@ -89,6 +92,10 @@ impl VTable for Not {
     fn is_fallible(&self, _instance: &Self::Instance) -> bool {
         false
     }
+
+    fn expr_v2(&self, view: &ExpressionView<Self>) -> VortexResult<Expression> {
+        ScalarFnExprExt::try_new_expr(&not::NotFn, EmptyOptions, view.children().clone())
+    }
 }
 
 /// Creates an expression that logically inverts boolean values.

vortex-array/src/expr/vtable.rs

@@ -8,6 +8,7 @@ use std::fmt::Display;
 use std::fmt::Formatter;
 use std::hash::Hash;
 use std::hash::Hasher;
+use std::ops::Deref;
 use std::sync::Arc;
 
 use arcref::ArcRef;
@@ -152,6 +153,13 @@ pub trait VTable: 'static + Sized + Send + Sync {
         _ = instance;
         true
     }
+
+    /// **For internal usage**. This will return an Expression that is part of the
+    /// expression -> new_expression migration.
+    #[doc(hidden)]
+    fn expr_v2(&self, expr: &ExpressionView<Self>) -> VortexResult<Expression> {
+        Ok(expr.deref().clone())
+    }
 }
 
 /// Arguments for expression execution.

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

@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ops::Not;
+
+use vortex_dtype::DType;
+use vortex_dtype::Nullability::NonNullable;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::Datum;
+use vortex_vector::ScalarOps;
+use vortex_vector::Vector;
+use vortex_vector::VectorOps;
+use vortex_vector::bool::BoolVector;
+
+use crate::expr::ChildName;
+use crate::expr::ExprId;
+use crate::expr::functions::ArgName;
+use crate::expr::functions::Arity;
+use crate::expr::functions::EmptyOptions;
+use crate::expr::functions::ExecutionArgs;
+use crate::expr::functions::VTable;
+
+pub struct IsNullFn;
+impl VTable for IsNullFn {
+    type Options = EmptyOptions;
+
+    fn id(&self) -> ExprId {
+        ExprId::new_ref("is_null")
+    }
+
+    fn arity(&self, _: &Self::Options) -> Arity {
+        Arity::Fixed(1)
+    }
+
+    fn arg_name(&self, _: &Self::Options, arg_idx: usize) -> ArgName {
+        match arg_idx {
+            0 => ChildName::from("input"),
+            _ => unreachable!("Invalid child index {} for IsNull expression", arg_idx),
+        }
+    }
+
+    fn return_dtype(&self, _options: &Self::Options, _arg_types: &[DType]) -> VortexResult<DType> {
+        Ok(DType::Bool(NonNullable))
+    }
+
+    fn execute(&self, _: &Self::Options, args: &ExecutionArgs) -> VortexResult<Datum> {
+        Ok(match args.input_datums(0) {
+            Datum::Scalar(sc) => Datum::Scalar(sc.is_invalid().into()),
+            Datum::Vector(vec) => Vector::Bool(BoolVector::new(
+                vec.validity().to_bit_buffer().not(),
+                Mask::AllTrue(vec.len()),
+            ))
+            .into(),
+        })
+    }
+}

vortex-array/src/scalar_fns/mod.rs

@@ -19,6 +19,8 @@ use crate::expr::Expression;
 use crate::expr::ScalarFnExprExt;
 
 pub mod cast;
+pub mod is_null;
+pub mod not;
 
 /// A collection of built-in scalar functions that can be applied to expressions or arrays.
 pub trait BuiltinScalarFns: Sized {

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

@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_compute::logical::LogicalNot;
+use vortex_dtype::DType;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_vector::Datum;
+use vortex_vector::Scalar;
+use vortex_vector::Vector;
+use vortex_vector::bool::BoolScalar;
+
+use crate::expr::ChildName;
+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::NullHandling;
+use crate::expr::functions::VTable;
+
+pub struct NotFn;
+impl VTable for NotFn {
+    type Options = EmptyOptions;
+
+    fn id(&self) -> FunctionId {
+        FunctionId::from("vortex.not")
+    }
+
+    fn arity(&self, _: &Self::Options) -> Arity {
+        Arity::Fixed(1)
+    }
+
+    fn null_handling(&self, _options: &Self::Options) -> NullHandling {
+        NullHandling::Propagate
+    }
+
+    fn arg_name(&self, _: &Self::Options, arg_idx: usize) -> ArgName {
+        match arg_idx {
+            0 => ChildName::from("input"),
+            _ => unreachable!("Invalid child index {} for Not expression", arg_idx),
+        }
+    }
+
+    fn return_dtype(&self, _options: &Self::Options, arg_types: &[DType]) -> VortexResult<DType> {
+        let child_dtype = &arg_types[0];
+        if !matches!(child_dtype, DType::Bool(_)) {
+            vortex_bail!(
+                "Not expression expects a boolean child, got: {}",
+                child_dtype
+            );
+        }
+        Ok(child_dtype.clone())
+    }
+
+    fn execute(&self, _: &Self::Options, args: &ExecutionArgs) -> VortexResult<Datum> {
+        Ok(match args.input_datums(0) {
+            Datum::Scalar(Scalar::Bool(sc)) => {
+                Datum::Scalar(BoolScalar::new(sc.value().map(|v| !v)).into())
+            }
+            Datum::Vector(Vector::Bool(vec)) => Datum::Vector(vec.clone().not().into()),
+            _ => unreachable!("Not expects a boolean"),
+        })
+    }
+}

vortex-vector/src/bool/scalar.rs

@@ -43,3 +43,9 @@ impl From<BoolScalar> for Scalar {
         Scalar::Bool(val)
     }
 }
+
+impl From<bool> for Scalar {
+    fn from(value: bool) -> Self {
+        BoolScalar::new(Some(value)).into()
+    }
+}