Lazy ScalarFn::validity

vortex-array/Cargo.toml

@@ -228,6 +228,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/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/scalar_fn/vtable/mod.rs

@@ -1,9 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 mod operations;
-mod validity;
-use std::fmt::Display;
-use std::fmt::Formatter;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::marker::PhantomData;
@@ -15,37 +12,40 @@
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
-use vortex_session::registry::CachedId;
 
 use crate::ArrayEq;
 use crate::ArrayHash;
 use crate::ArrayRef;
 use crate::EqMode;
+use crate::IntoArray;
 use crate::array::Array;
 use crate::array::ArrayId;
 use crate::array::ArrayParts;
 use crate::array::ArrayView;
 use crate::array::VTable;
+use crate::array::ValidityVTable;
 use crate::array::with_empty_buffers;
+use crate::arrays::StructArray;
 use crate::arrays::scalar_fn::array::ScalarFnArrayExt;
 use crate::arrays::scalar_fn::array::ScalarFnData;
 use crate::arrays::scalar_fn::rules::PARENT_RULES;
 use crate::arrays::scalar_fn::rules::RULES;
 use crate::buffer::BufferHandle;
 use crate::dtype::DType;
+use crate::dtype::FieldName;
 use crate::executor::ExecutionCtx;
 use crate::executor::ExecutionResult;
 use crate::expr::Expression;
-use crate::expr::display::ExprDisplay;
+use crate::expr::get_item;
+use crate::expr::is_not_null;
+use crate::expr::lit;
+use crate::expr::root;
 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::ScalarFnId;
-use crate::scalar_fn::ScalarFnVTableExt;
 use crate::scalar_fn::VecExecutionArgs;
 use crate::serde::ArrayChildren;
+use crate::validity::Validity;
 
 /// A [`ScalarFn`]-encoded Vortex array.
 pub type ScalarFnArray = Array<ScalarFn>;
@@ -244,79 +244,55 @@
     }
 }
 
-// Used only in this method to allow constrained using of Expression evaluate.
-#[derive(Clone)]
-struct ArrayExpr;
-
-#[derive(Clone, Debug)]
-struct FakeEq<T>(T);
-
-impl<T> PartialEq<Self> for FakeEq<T> {
-    fn eq(&self, _other: &Self) -> bool {
-        false
-    }
-}
-
-impl<T> Eq for FakeEq<T> {}
-
-impl<T> Hash for FakeEq<T> {
-    fn hash<H: Hasher>(&self, _state: &mut H) {}
-}
-
-impl Display for FakeEq<ArrayRef> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.0.encoding_id())
-    }
-}
-
-impl scalar_fn::ScalarFnVTable for ArrayExpr {
-    type Options = FakeEq<ArrayRef>;
-
-    fn id(&self) -> ScalarFnId {
-        static ID: CachedId = CachedId::new("vortex.array");
-        *ID
-    }
-
-    fn arity(&self, _options: &Self::Options) -> Arity {
-        Arity::Exact(0)
-    }
-
-    fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName {
-        todo!()
-    }
-
-    fn fmt_sql(
-        &self,
-        options: &Self::Options,
-        _expr: &dyn ExprDisplay,
-        f: &mut Formatter<'_>,
-    ) -> std::fmt::Result {
-        write!(f, "{}", options.0.encoding_id())
-    }
-
-    fn return_dtype(&self, options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
-        Ok(options.0.dtype().clone())
-    }
-
-    fn execute(
-        &self,
-        options: &Self::Options,
-        _args: &dyn ExecutionArgs,
-        ctx: &mut ExecutionCtx,
-    ) -> VortexResult<ArrayRef> {
-        crate::Executable::execute(options.0.clone(), ctx)
-    }
-
-    fn validity(
-        &self,
-        options: &Self::Options,
-        _expression: &Expression,
-    ) -> VortexResult<Option<Expression>> {
-        let validity_array = options.0.validity()?.to_array(options.0.len());
-        Ok(Some(ArrayExpr.new_expr(FakeEq(validity_array), [])))
-    }
-
-    fn is_strict(&self, _options: &Self::Options) -> bool {
-        true
+impl ValidityVTable<ScalarFn> for ScalarFn {
+    fn validity(view: ArrayView<'_, ScalarFn>) -> VortexResult<Validity> {
+        // We want to defer execution of the underlying array. The naïve
+        // solution for this is to build an all true array and then .apply() an
+        // Expression referencing parts of root(). This doesn't work because
+        // in this Expression's evaluation root() is replaced by the original
+        // array which leads to non-terminating recursion, a stack overflow. So
+        // we build a Struct array and give the caller (which overrides
+        // ScalarFn valididy) the ability to reference children with get_item.
+        // In ScalarFn's overriden validity "expr.child(i)" then translates to
+        // "view.get_item(i)" which doesn't produce recursion since get_item
+        // references part of the original array as opposed to root().
+        let child_count = view.child_count();
+        let names = (0..child_count)
+            .map(|i| FieldName::from(i.to_string().as_str()))
+            .collect();
+        let fields = view.children();
+
+        let getters: Vec<_> = view
+            .children()
+            .into_iter()
+            .enumerate()
+            .map(|(i, child)| {
+                if let Some(scalar) = child.as_constant() {
+                    lit(scalar)
+                } else {
+                    get_item(i.to_string(), root())
+                }
+            })
+            .collect();
+
+        let struct_array = StructArray::new(names, fields, view.len(), Validity::NonNullable);
+
+        let scalar_fn = view.scalar_fn();
+        let expr = Expression::try_new(scalar_fn.clone(), getters)?;
+        let expr = scalar_fn
+            .validity(&expr)?
+            // However, there is another possible stack overflow if validity()
+            // isn't overriden. The naïve solution is to do is_not_null(expr)
+            // which is is_not_null(ScalarFn(get_item(...))).
+            // Inner ScalarFn(get_item)'s row request will call validity() back
+            // which will instantiate is_not_null(F( original is_not_null )).
+            //
+            // So, to break this recursion, we need to tweak array probing for
+            // ScalarFn, see execute_scalar in probe/array.rs and in
+            // probe/repeated.rs
+            .unwrap_or_else(|| is_not_null(expr.clone()));
+
+        let array = struct_array.into_array().apply(&expr)?;
+        Ok(Validity::Array(array))
     }
 }