ScalarFn validity

vortex-array/src/array/erased.rs

@@ -41,6 +41,7 @@ use crate::arrays::DictArray;
 use crate::arrays::FilterArray;
 use crate::arrays::Null;
 use crate::arrays::Primitive;
+use crate::arrays::ScalarFn;
 use crate::arrays::SliceArray;
 use crate::arrays::VarBin;
 use crate::arrays::VarBinView;
@@ -276,7 +277,7 @@ impl ArrayRef {
     /// Execute the array to extract a scalar at the given index.
     pub fn execute_scalar(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
         vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len());
-        if self.dtype().is_nullable() && self.is_invalid(index, ctx)? {
+        if self.dtype().is_nullable() && !self.is::<ScalarFn>() && self.is_invalid(index, ctx)? {
             return Ok(Scalar::null(self.dtype().clone()));
         }
         let scalar = self.0.data.execute_scalar(self, index, ctx)?;

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

@@ -8,6 +8,7 @@ use std::hash::Hash;
 use std::hash::Hasher;
 use std::marker::PhantomData;
 use std::ops::Deref;
+use std::sync::Arc;
 
 use itertools::Itertools;
 use vortex_error::VortexResult;
@@ -37,13 +38,18 @@ use crate::dtype::DType;
 use crate::executor::ExecutionCtx;
 use crate::executor::ExecutionResult;
 use crate::expr::Expression;
+use crate::expr::lit;
 use crate::matcher::Matcher;
 use crate::scalar_fn;
 use crate::scalar_fn::Arity;
 use crate::scalar_fn::ChildName;
 use crate::scalar_fn::ExecutionArgs;
+use crate::scalar_fn::ReduceCtx;
+use crate::scalar_fn::ReduceNode;
+use crate::scalar_fn::ReduceNodeRef;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::ScalarFnVTableExt;
+use crate::scalar_fn::TypedScalarFnInstance;
 use crate::scalar_fn::VecExecutionArgs;
 use crate::serde::ArrayChildren;
 
@@ -175,7 +181,7 @@ pub trait ScalarFnFactoryExt: scalar_fn::ScalarFnVTable {
         options: Self::Options,
         children: impl Into<Vec<ArrayRef>>,
     ) -> VortexResult<ArrayRef> {
-        let scalar_fn = scalar_fn::TypedScalarFnInstance::new(self.clone(), options).erased();
+        let scalar_fn = TypedScalarFnInstance::new(self.clone(), options).erased();
 
         let children = children.into();
         vortex_ensure!(
@@ -201,6 +207,24 @@ pub trait ScalarFnFactoryExt: scalar_fn::ScalarFnVTable {
 }
 impl<V: scalar_fn::ScalarFnVTable> ScalarFnFactoryExt for V {}
 
+pub(crate) fn scalar_fn_array_expr(array: ArrayView<'_, ScalarFn>) -> VortexResult<Expression> {
+    let inputs: Vec<_> = array
+        .iter_children()
+        .map(|child| {
+            if let Some(scalar) = child.as_constant() {
+                return Ok(lit(scalar));
+            }
+
+            Expression::try_new(
+                TypedScalarFnInstance::new(ArrayExpr, FakeEq(child.clone())).erased(),
+                [],
+            )
+        })
+        .collect::<VortexResult<_>>()?;
+
+    Expression::try_new(array.scalar_fn().clone(), inputs)
+}
+
 /// A matcher that matches any scalar function expression.
 #[derive(Debug)]
 pub struct AnyScalarFn;
@@ -320,12 +344,21 @@ impl scalar_fn::ScalarFnVTable for ArrayExpr {
         crate::Executable::execute(options.0.clone(), ctx)
     }
 
+    fn reduce(
+        &self,
+        options: &Self::Options,
+        _node: &dyn ReduceNode,
+        _ctx: &dyn ReduceCtx,
+    ) -> VortexResult<Option<ReduceNodeRef>> {
+        Ok(Some(Arc::new(options.0.clone())))
+    }
+
     fn validity(
         &self,
         options: &Self::Options,
         _expression: &Expression,
-    ) -> VortexResult<Option<Expression>> {
+    ) -> VortexResult<Expression> {
         let validity_array = options.0.validity()?.to_array(options.0.len());
-        Ok(Some(ArrayExpr.new_expr(FakeEq(validity_array), [])))
+        Ok(ArrayExpr.new_expr(FakeEq(validity_array), []))
     }
 }

vortex-array/src/arrays/scalar_fn/vtable/operations.rs

@@ -65,6 +65,7 @@ mod tests {
     use crate::arrays::BoolArray;
     use crate::arrays::PrimitiveArray;
     use crate::arrays::ScalarFnArray;
+    use crate::arrays::StructArray;
     use crate::arrays::scalar_fn::ScalarFnArrayExt;
     use crate::assert_arrays_eq;
     use crate::scalar::Scalar;
@@ -172,6 +173,29 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn scalar_fn_scalar_at_handles_value_derived_validity() -> VortexResult<()> {
+        let child = StructArray::from_fields(&[(
+            "a",
+            PrimitiveArray::from_option_iter([Some(1i32), None]).into_array(),
+        )])?
+        .into_array();
+        let expr = crate::expr::get_item("a", crate::expr::root());
+        let array = ScalarFnArray::try_new(expr.scalar_fn().clone(), vec![child])?.into_array();
+
+        assert_eq!(
+            array.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?,
+            Scalar::primitive(1i32, true.into())
+        );
+        assert!(
+            array
+                .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())?
+                .is_null()
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn test_scalar_fn_comparison() -> VortexResult<()> {
         let lhs = buffer![1i32, 5, 3].into_array();

vortex-array/src/arrays/scalar_fn/vtable/validity.rs

@@ -3,72 +3,57 @@
 
 use vortex_error::VortexResult;
 
-use crate::ArrayRef;
 use crate::IntoArray;
-use crate::LEGACY_SESSION;
-use crate::VortexSessionExecute;
 use crate::array::ArrayView;
 use crate::array::ValidityVTable;
-use crate::arrays::scalar_fn::ScalarFnArrayExt;
-use crate::arrays::scalar_fn::vtable::ArrayExpr;
-use crate::arrays::scalar_fn::vtable::FakeEq;
+use crate::arrays::ConstantArray;
 use crate::arrays::scalar_fn::vtable::ScalarFn;
-use crate::expr::Expression;
-use crate::expr::lit;
-use crate::scalar_fn::TypedScalarFnInstance;
-use crate::scalar_fn::VecExecutionArgs;
-use crate::scalar_fn::fns::literal::Literal;
-use crate::scalar_fn::fns::root::Root;
+use crate::arrays::scalar_fn::vtable::scalar_fn_array_expr;
 use crate::validity::Validity;
 
-/// Execute an expression tree recursively.
-///
-/// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals.
-fn execute_expr(expr: &Expression, row_count: usize) -> VortexResult<ArrayRef> {
-    let mut ctx = LEGACY_SESSION.create_execution_ctx();
-
-    // Handle Root expression - this should not happen in validity expressions
-    if expr.is::<Root>() {
-        vortex_error::vortex_bail!("Root expression cannot be executed in validity context");
-    }
-
-    // Handle Literal expression - create a constant array
-    if expr.is::<Literal>() {
-        let scalar = expr.as_::<Literal>();
-        return Ok(crate::arrays::ConstantArray::new(scalar.clone(), row_count).into_array());
-    }
-
-    // Recursively execute child expressions to get input arrays
-    let inputs: Vec<ArrayRef> = expr
-        .children()
-        .iter()
-        .map(|child| execute_expr(child, row_count))
-        .collect::<VortexResult<_>>()?;
-
-    let args = VecExecutionArgs::new(inputs, row_count);
-
-    Ok(expr.scalar_fn().execute(&args, &mut ctx)?.into_array())
-}
-
 impl ValidityVTable<ScalarFn> for ScalarFn {
     fn validity(array: ArrayView<'_, ScalarFn>) -> VortexResult<Validity> {
-        let inputs: Vec<_> = array
-            .iter_children()
-            .map(|child| {
-                if let Some(scalar) = child.as_constant() {
-                    return Ok(lit(scalar));
-                }
-                Expression::try_new(
-                    TypedScalarFnInstance::new(ArrayExpr, FakeEq(child.clone())).erased(),
-                    [],
-                )
-            })
-            .collect::<VortexResult<_>>()?;
-
-        let expr = Expression::try_new(array.scalar_fn().clone(), inputs)?;
-        let validity_expr = array.scalar_fn().validity(&expr)?;
+        let expr = scalar_fn_array_expr(array)?.validity()?;
+        let input = ConstantArray::new(true, array.len()).into_array();
+        Ok(Validity::Array(input.apply(&expr)?))
+    }
+}
 
-        // Execute the validity expression. All leaves are ArrayExpr nodes.
-        Ok(Validity::Array(execute_expr(&validity_expr, array.len())?))
+#[cfg(test)]
+mod tests {
+    use vortex_error::VortexResult;
+    use vortex_mask::Mask;
+
+    use crate::IntoArray;
+    use crate::LEGACY_SESSION;
+    use crate::VortexSessionExecute;
+    use crate::arrays::BoolArray;
+    use crate::arrays::ScalarFn;
+    use crate::arrays::scalar_fn::ScalarFnArrayExt;
+    use crate::arrays::scalar_fn::vtable::ScalarFnFactoryExt;
+    use crate::scalar_fn::fns::binary::Binary;
+    use crate::scalar_fn::fns::operators::Operator;
+    use crate::validity::Validity;
+
+    #[test]
+    fn scalar_fn_validity_stays_lazy() -> VortexResult<()> {
+        let lhs = BoolArray::from_iter([Some(true), None, Some(false)]).into_array();
+        let rhs = BoolArray::from_iter([Some(true), Some(false), None]).into_array();
+        let predicate = Binary.try_new_array(lhs.len(), Operator::And, [lhs, rhs])?;
+
+        let Validity::Array(validity_array) = predicate.validity()? else {
+            panic!("scalar function validity should be represented as an array");
+        };
+
+        let validity_scalar_fn = validity_array
+            .as_opt::<ScalarFn>()
+            .expect("validity should remain a lazy scalar function array");
+        assert!(validity_scalar_fn.scalar_fn().is::<Binary>());
+
+        let validity_mask =
+            validity_array.execute::<Mask>(&mut LEGACY_SESSION.create_execution_ctx())?;
+        assert!(validity_mask.all_true());
+
+        Ok(())
     }
 }

vortex-array/src/scalar_fn/erased.rs

@@ -20,16 +20,13 @@ use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::dtype::DType;
 use crate::expr::Expression;
-use crate::scalar_fn::EmptyOptions;
 use crate::scalar_fn::ExecutionArgs;
 use crate::scalar_fn::ReduceCtx;
 use crate::scalar_fn::ReduceNode;
 use crate::scalar_fn::ReduceNodeRef;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::ScalarFnVTable;
-use crate::scalar_fn::ScalarFnVTableExt;
 use crate::scalar_fn::SimplifyCtx;
-use crate::scalar_fn::fns::is_not_null::IsNotNull;
 use crate::scalar_fn::options::ScalarFnOptions;
 use crate::scalar_fn::signature::ScalarFnSignature;
 use crate::scalar_fn::typed::DynScalarFn;
@@ -132,10 +129,7 @@ impl ScalarFnRef {
 
     /// Transforms the expression into one representing the validity of this expression.
     pub fn validity(&self, expr: &Expression) -> VortexResult<Expression> {
-        Ok(self.0.validity(expr)?.unwrap_or_else(|| {
-            // TODO(ngates): make validity a mandatory method on VTable to avoid this fallback.
-            IsNotNull.new_expr(EmptyOptions, [expr.clone()])
-        }))
+        self.0.validity(expr)
     }
 
     /// Execute the expression given the input arguments.

vortex-array/src/scalar_fn/fns/between/mod.rs

@@ -299,11 +299,11 @@ impl ScalarFnVTable for Between {
         &self,
         _options: &Self::Options,
         expression: &Expression,
-    ) -> VortexResult<Option<Expression>> {
+    ) -> VortexResult<Expression> {
         let arr = expression.child(0).validity()?;
         let lower = expression.child(1).validity()?;
         let upper = expression.child(2).validity()?;
-        Ok(Some(and(and(arr, lower), upper)))
+        Ok(and(and(arr, lower), upper))
     }
 
     fn is_null_sensitive(&self, _instance: &Self::Options) -> bool {