Lazy ScalarFn::validity

vortex-array/Cargo.toml

@@ -236,6 +236,10 @@ harness = false
 name = "validity_is_valid"
 harness = false
 
+[[bench]]
+name = "scalar_fn_probe"
+harness = false
+
 [[bench]]
 name = "dict_unreferenced_mask"
 harness = false

vortex-array/benches/scalar_fn_probe.rs

@@ -0,0 +1,124 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![expect(clippy::unwrap_used)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use divan::black_box;
+use rand::RngExt;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::arrays::ScalarFnArray;
+use vortex_array::scalar_fn::TypedScalarFnInstance;
+use vortex_array::scalar_fn::fns::binary::Binary;
+use vortex_array::scalar_fn::fns::operators::Operator;
+use vortex_session::VortexSession;
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+const ARRAY_SIZE: usize = 100_000;
+const NUM_ACCESSES: usize = 50;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
+
+/// evaluating ADD's validity is cheaper than evaluating ADD
+fn binary_add() -> ArrayRef {
+    let lhs =
+        PrimitiveArray::from_option_iter((0..ARRAY_SIZE).map(|i| (i % 7 != 0).then_some(i as i64)))
+            .into_array();
+    let rhs = PrimitiveArray::from_iter((0..ARRAY_SIZE).map(|i| i as i64)).into_array();
+    let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::Add).erased();
+    ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])
+        .unwrap()
+        .into_array()
+}
+
+/// evaluating AND's validity is equal to evaluating the AND due to
+/// Kleene semantics
+fn binary_and() -> ArrayRef {
+    let lhs = BoolArray::from_iter((0..ARRAY_SIZE).map(|i| (i % 7 != 0).then_some(i % 2 == 0)))
+        .into_array();
+    let rhs = BoolArray::from_iter((0..ARRAY_SIZE).map(|i| i % 2 == 0)).into_array();
+    let scalar_fn = TypedScalarFnInstance::new(Binary, Operator::And).erased();
+    ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])
+        .unwrap()
+        .into_array()
+}
+
+fn indices() -> Vec<usize> {
+    let mut rng = StdRng::seed_from_u64(0);
+    (0..NUM_ACCESSES)
+        .map(|_| rng.random_range(0..ARRAY_SIZE))
+        .collect()
+}
+
+#[divan::bench(args = [binary_and(), binary_add()])]
+fn probe_scalar_fn_once(bencher: Bencher, array: &ArrayRef) {
+    let indices = indices();
+    bencher
+        .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx()))
+        .bench_refs(|(array, indices, ctx)| {
+            for &index in indices.iter() {
+                black_box(array.probe().execute_scalar(index, ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = [binary_and(), binary_add()])]
+fn probe_scalar_fn_repeated(bencher: Bencher, array: &ArrayRef) {
+    let indices = indices();
+    bencher
+        .with_inputs(|| {
+            (
+                array.repeated_probe(),
+                &indices,
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(probe, indices, ctx)| {
+            for &index in indices.iter() {
+                black_box(probe.execute_scalar(index, ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = [binary_and(), binary_add()])]
+fn probe_scalar_fn_valid_once(bencher: Bencher, array: &ArrayRef) {
+    let indices = indices();
+    bencher
+        .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx()))
+        .bench_refs(|(array, indices, ctx)| {
+            for &index in indices.iter() {
+                black_box(array.probe().execute_is_valid(index, ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = [binary_and(), binary_add()])]
+fn probe_scalar_fn_valid_repeated(bencher: Bencher, array: &ArrayRef) {
+    let indices = indices();
+    bencher
+        .with_inputs(|| {
+            (
+                array.repeated_probe(),
+                &indices,
+                SESSION.create_execution_ctx(),
+            )
+        })
+        .bench_refs(|(probe, indices, ctx)| {
+            for &index in indices.iter() {
+                black_box(probe.execute_is_invalid(index, ctx).unwrap());
+            }
+        });
+}

vortex-array/src/array/mod.rs

@@ -16,6 +16,7 @@ use vortex_error::vortex_panic;
 use vortex_session::registry::Id;
 
 use crate::ExecutionCtx;
+use crate::arrays::Constant;
 use crate::buffer::BufferHandle;
 use crate::builders::ArrayBuilder;
 use crate::dtype::DType;
@@ -427,12 +428,15 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
             this.encoding_id(),
             reduced.encoding_id()
         );
-        vortex_ensure!(
-            reduced.dtype() == this.dtype(),
-            "Reduced array dtype mismatch from {} to {}",
-            this.encoding_id(),
-            reduced.encoding_id()
-        );
+        // Array can be reduced to constant during optimize()
+        if reduced.encoding_id() != VTable::id(&Constant) {
+            vortex_ensure!(
+                reduced.dtype() == this.dtype(),
+                "Reduced array dtype mismatch from {} to {}",
+                this.encoding_id(),
+                reduced.encoding_id()
+            );
+        }
         Ok(Some(reduced))
     }
 

vortex-array/src/array/probe/array.rs

@@ -15,6 +15,7 @@ use crate::array::probe::RepeatedArrayProbe;
 use crate::array::probe::RepeatedState;
 use crate::array::probe::repeated::child_probe;
 use crate::arrays::Primitive;
+use crate::arrays::ScalarFn;
 use crate::scalar::Scalar;
 use crate::vtable::OperationsVTable;
 
@@ -82,7 +83,17 @@ fn execute_scalar_once(
     index: usize,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<Scalar> {
-    if !execute_is_valid_once(array, index, ctx)? {
+    // ScalarFn's validity is lazy, and for some functions evaluating
+    // validity is equal to evaluating the function. For such functions
+    // validity() is is_not_null(original array). So we get the chain:
+    // execute_is_valid_once -> array.validity() ->
+    // execute_is_valid -> execute_scalar (mask) ->
+    // mask.probe_scalar_once -> scalar_at -> array.execute_scalar, and as
+    // "array" is the original array, we get infinite recursion.
+    //
+    // For these functions probe_scalar_once gets the nullable scalar anyway.
+    // See also execute_scalar in probe/repeated.rs
+    if !array.is::<ScalarFn>() && !execute_is_valid_once(array, index, ctx)? {
         return Ok(Scalar::null(array.dtype().clone()));
     }
     check_dtype(

vortex-array/src/array/probe/repeated.rs

@@ -12,6 +12,7 @@ use crate::array::probe::ArrayProbe;
 use crate::array::probe::array::check_bounds;
 use crate::array::probe::array::check_dtype;
 use crate::array::probe::array::child_of;
+use crate::arrays::ScalarFn;
 use crate::scalar::Scalar;
 use crate::validity::Validity;
 
@@ -55,7 +56,18 @@ impl RepeatedArrayProbe {
 
     /// Read the scalar at `index`, including its nullness, reusing retained preparation.
     pub fn execute_scalar(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
-        if !self.execute_is_valid(index, ctx)? {
+        // ScalarFn's validity is lazy, and for some functions evaluating
+        // validity is equal to evaluating the function. For such functions
+        // validity() is is_not_null(original array). So we get the chain:
+        // execute_scalar -> array.validity() ->
+        // execute_is_valid -> execute_scalar (mask) ->
+        // mask.probe_scalar_once -> scalar_at -> array.execute_scalar, and as
+        // "array" is the original array, we get infinite recursion.
+        //
+        // For these functions probe_scalar_once gets the nullable scalar anyway.
+        //
+        // See also execute_scala_once in probe/array.rs
+        if !self.array.is::<ScalarFn>() && !self.execute_is_valid(index, ctx)? {
             return Ok(Scalar::null(self.array.dtype().clone()));
         }
         let result =

vortex-array/src/arrays/dict/compute/rules.rs

@@ -369,8 +369,8 @@ mod tests {
         let expr = Binary::try_new(dict, const_false, Operator::And)?.into_array();
 
         let mut ctx = array_session().create_execution_ctx();
-        // Kleene AND: null AND false == false, so all three rows must be valid `false`.
-        let expected = expr.clone().execute::<BoolArray>(&mut ctx)?.into_array();
+        // Kleene AND: null AND false == false
+        let expected = BoolArray::from_iter([false, false, false]).into_array();
         let optimized = expr.optimize()?;
         assert_arrays_eq!(optimized, expected, &mut ctx);
         Ok(())
@@ -386,8 +386,8 @@ mod tests {
         let expr = Binary::try_new(dict, const_true, Operator::Or)?.into_array();
 
         let mut ctx = array_session().create_execution_ctx();
-        // Kleene OR: null OR true == true, so all three rows must be valid `true`.
-        let expected = expr.clone().execute::<BoolArray>(&mut ctx)?.into_array();
+        // Kleene OR: null OR true == true
+        let expected = BoolArray::from_iter([true, true, true]).into_array();
         let optimized = expr.optimize()?;
         assert_arrays_eq!(optimized, expected, &mut ctx);
         Ok(())

vortex-array/src/arrays/scalar_fn/rules.rs

@@ -29,8 +29,10 @@ use crate::validity::Validity;
 
 pub(super) const RULES: ReduceRuleSet<ScalarFn> = ReduceRuleSet::new(&[
     &ScalarFnPackToStructRule,
-    &ScalarFnAbstractReduceRule,
     &IsNullReduceRule,
+    // Ordering is important. ScalarFn::reduce() must be called after all other
+    // optimizations
+    &ScalarFnAbstractReduceRule,
 ]);
 
 pub(super) const PARENT_RULES: ParentRuleSet<ScalarFn> = ParentRuleSet::new(&[
@@ -86,7 +88,12 @@ impl ArrayReduceRule<ScalarFn> for IsNullReduceRule {
             Validity::AllInvalid => ConstantArray::new(is_null, view.len()).into_array(),
             Validity::Array(array) => {
                 if is_null {
+<<<<<<< HEAD
                     array.not()?
+||||||| parent of 63300294f (reduce is_null(array) to array's validity)
+=======
+                    array.apply(&not(root()))?
+>>>>>>> 63300294f (reduce is_null(array) to array's validity)
                 } else {
                     array
                 }
@@ -203,6 +210,7 @@ mod tests {
     use crate::expr::cast;
     use crate::expr::is_not_null;
     use crate::expr::is_null;
+    use crate::expr::not;
     use crate::expr::root;
     use crate::optimizer::rules::ArrayParentReduceRule;
     use crate::scalar::Scalar;
@@ -332,8 +340,22 @@ mod tests {
         let buffer = buffer![1, 2, 3];
         let nullable = PrimitiveArray::new(buffer, Validity::Array(validity.clone())).into_array();
 
+<<<<<<< HEAD
         assert_arrays_eq!(nullable.clone().apply(&is_not_null(root()))?, validity, ctx);
         assert_arrays_eq!(nullable.apply(&is_null(root()))?, validity.not()?, ctx);
+||||||| parent of 63300294f (reduce is_null(array) to array's validity)
+=======
+        assert_arrays_eq!(
+            nullable.clone().apply(&is_not_null(root()))?,
+            validity.clone(),
+            ctx
+        );
+        assert_arrays_eq!(
+            nullable.apply(&is_null(root()))?,
+            validity.apply(&not(root()))?,
+            ctx
+        );
+>>>>>>> 63300294f (reduce is_null(array) to array's validity)
 
         Ok(())
     }