From 3d9515439a896c5cc966d8a1c4a090e63f243fdd Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 10 Jun 2026 13:21:21 +0000 Subject: [PATCH] Make ScalarFn array validity lazy when the function defines a validity expression Previously ValidityVTable always eagerly executed the validity expression via the legacy session. Now, when the scalar function provides a validity expression over its inputs, the expression is converted into a lazy ScalarFn array DAG instead: Literal nodes become ConstantArrays, ArrayExpr leaves unwrap to the child arrays they hold, and interior nodes become lazy ScalarFn arrays. Constant results are folded back into AllValid/AllInvalid via child_to_validity. Functions that do not define a validity expression (e.g. Kleene logic and/or, where validity depends on the computed values) keep the eager path. The erased fallback for these is is_not_null over the expression itself, so a lazy representation would be self-referential: resolving the validity of the inner node spawns another is_not_null DAG, which recurses without ever shrinking (this manifested as a stack overflow in element-wise execution paths). ScalarFnRef::validity_opt is added to expose whether a function defines its own validity expression. https://claude.ai/code/session_01VPQ7dfZtijfrsjAipwXvEj Signed-off-by: Joe Isaacs --- .../src/arrays/scalar_fn/vtable/validity.rs | 71 ++++++++++++++++--- vortex-array/src/scalar_fn/erased.rs | 12 +++- 2 files changed, 73 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 39ab8c2c466..89068c4f60d 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs @@ -8,13 +8,16 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::Array; use crate::array::ArrayView; use crate::array::ValidityVTable; +use crate::array::child_to_validity; use crate::arrays::ConstantArray; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::arrays::scalar_fn::vtable::ArrayExpr; use crate::arrays::scalar_fn::vtable::FakeEq; use crate::arrays::scalar_fn::vtable::ScalarFn; +use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::lit; use crate::legacy_session; @@ -24,6 +27,39 @@ use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::root::Root; use crate::validity::Validity; +/// Convert an expression tree into a lazy array DAG without executing it. +/// +/// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals. +fn expr_to_lazy_array(expr: &Expression, row_count: usize) -> VortexResult { + // Handle Root expression - this should not happen in validity expressions + if expr.is::() { + vortex_bail!("Root expression cannot be converted in validity context"); + } + + // Handle Literal expression - create a constant array + if expr.is::() { + let scalar = expr.as_::(); + return Ok(ConstantArray::new(scalar.clone(), row_count).into_array()); + } + + // Handle ArrayExpr leaves - unwrap the array they hold + if expr.is::() { + return Ok(expr.as_::().0.clone()); + } + + // Recursively convert child expressions into lazy input arrays + let children: Vec = expr + .children() + .iter() + .map(|child| expr_to_lazy_array(child, row_count)) + .collect::>()?; + + Ok( + Array::::try_new_with_len(expr.scalar_fn().clone(), children, row_count)? + .into_array(), + ) +} + /// Execute an expression tree recursively. /// /// This assumes all leaf expressions are either ArrayExpr (wrapping actual arrays) or Literals. @@ -71,15 +107,32 @@ impl ValidityVTable for ScalarFn { .collect::>()?; let expr = Expression::try_new(array.scalar_fn().clone(), inputs)?; - let validity_expr = array.scalar_fn().validity(&expr)?; - #[allow(clippy::disallowed_methods)] - let ctx = &mut legacy_session().create_execution_ctx(); - // Execute the validity expression. All leaves are ArrayExpr nodes. - Ok(Validity::Array(execute_expr( - &validity_expr, - array.len(), - ctx, - )?)) + match array.scalar_fn().validity_opt(&expr)? { + Some(validity_expr) => { + // The function defines its validity as an expression over its inputs, so we can + // represent it as a lazy array DAG without executing anything. If the expression + // is already a constant it is folded back into AllValid/AllInvalid. + let validity_array = expr_to_lazy_array(&validity_expr, array.len())?; + Ok(child_to_validity( + Some(&validity_array), + Nullability::Nullable, + )) + } + None => { + // The function's validity can only be determined by executing the function + // itself (e.g. Kleene logic and/or). Representing that lazily would create a + // self-referential array (is_not_null over this very expression), so execute it + // eagerly instead. + let validity_expr = array.scalar_fn().validity(&expr)?; + #[allow(clippy::disallowed_methods)] + let ctx = &mut legacy_session().create_execution_ctx(); + Ok(Validity::Array(execute_expr( + &validity_expr, + array.len(), + ctx, + )?)) + } + } } } diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index 6e0011c297a..96879e6c3e7 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -132,12 +132,22 @@ impl ScalarFnRef { /// Transforms the expression into one representing the validity of this expression. pub fn validity(&self, expr: &Expression) -> VortexResult { - Ok(self.0.validity(expr)?.unwrap_or_else(|| { + Ok(self.validity_opt(expr)?.unwrap_or_else(|| { // TODO(ngates): make validity a mandatory method on VTable to avoid this fallback. IsNotNull.new_expr(EmptyOptions, [expr.clone()]) })) } + /// Transforms the expression into one representing the validity of this expression, + /// returning `None` if the function does not define a validity expression. + /// + /// When `None` is returned, the validity can only be determined by executing the + /// expression itself (e.g. Kleene logic `and`/`or`), and [`Self::validity`] falls back to + /// `is_not_null` over the expression. + pub fn validity_opt(&self, expr: &Expression) -> VortexResult> { + self.0.validity(expr) + } + /// Execute the expression given the input arguments. pub fn execute( &self,