Plan and execute common RowFn batches in one dispatch

vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs

@@ -9,22 +9,27 @@ use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
+use crate::validity::Validity;
 
-impl RowFnExecutionArgs {
+impl RowFnExecutionArgs<'_> {
     /// Execute all-constant inputs by evaluating one row and broadcasting the validated result.
     pub(super) fn execute_all_constant(
         &self,
-        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        kernel: impl FnOnce(
+            BorrowedRowFnArgs<'_>,
+            &mut Option<Validity>,
+            &mut ExecutionCtx,
+        ) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
         let result = self.validate_kernel_output(
-            kernel(self.execution_args(&self.inputs, 1), ctx)?,
+            kernel(self.execution_args(&self.batch.inputs, 1), &mut None, ctx)?,
             1,
             ctx,
         )?;
         let result = self.finalize_output(result, 1)?;
         let scalar = result.execute_scalar(0, ctx)?;
 
-        Ok(ConstantArray::new(scalar, self.row_count).into_array())
+        Ok(ConstantArray::new(scalar, self.batch.row_count).into_array())
     }
 }

vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs

@@ -9,19 +9,39 @@ use vortex_mask::MaskValuesRef;
 use super::super::RowFnExecutionArgs;
 use super::super::args::BorrowedRowFnArgs;
 use crate::ArrayRef;
+use crate::Canonical;
 use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::arrays::masked::mask_validity_canonical;
 use crate::builtins::ArrayBuiltins;
+use crate::canonical::AnyCanonical;
 use crate::scalar_fn::unstable::row::execute::DenseAttempt;
 use crate::validity::Validity;
 
-impl RowFnExecutionArgs {
+impl RowFnExecutionArgs<'_> {
     /// Run every stored payload, then attach the input validity without materializing its mask.
     pub(super) fn execute_dense(
         &self,
-        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        kernel: impl FnOnce(
+            BorrowedRowFnArgs<'_>,
+            &mut Option<Validity>,
+            &mut ExecutionCtx,
+        ) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?;
+        // Offer a partial validity to the kernel. A collector that takes it builds the output with
+        // that validity attached, so the output needs neither a second construction nor an
+        // all-valid check.
+        let mut validity =
+            matches!(self.batch.validity, Validity::Array(_)).then(|| self.batch.validity.clone());
+        let values = kernel(
+            self.execution_args(&self.batch.inputs, self.batch.row_count),
+            &mut validity,
+            ctx,
+        )?;
+        if validity.is_none() && matches!(self.batch.validity, Validity::Array(_)) {
+            return self.finalize_output(values, self.batch.row_count);
+        }
 
         self.finalize_dense_output(values, ctx)
     }
@@ -49,8 +69,10 @@ impl RowFnExecutionArgs {
         ) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let attempt =
-            execute_dense_attempt(self.execution_args(&self.inputs, self.row_count), ctx)?;
+        let attempt = execute_dense_attempt(
+            self.execution_args(&self.batch.inputs, self.batch.row_count),
+            ctx,
+        )?;
 
         match attempt {
             DenseAttempt::Values(values) => self.finalize_dense_output(values, ctx),
@@ -75,11 +97,15 @@ impl RowFnExecutionArgs {
         ) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let valid_rows = match self.validity.execute_mask(self.row_count, ctx)? {
+        let valid_rows = match self
+            .batch
+            .validity
+            .execute_mask(self.batch.row_count, ctx)?
+        {
             // An array-backed validity can materialize to all valid even though the cheap checks
-            // in `RowFnExecutionArgs::execute` could not prove that. The deferred error therefore
-            // came from an observable row and remains terminal. An empty mask is both all-valid
-            // and all-null, so preserve the all-valid behavior.
+            // in `RowFnExecutionArgs::try_execute_without_retry` could not prove that. The
+            // deferred error therefore came from an observable row and remains terminal. An empty
+            // mask is both all-valid and all-null, so preserve the all-valid behavior.
             Mask::AllTrue(_) | Mask::AllFalse(0) => return Err(deferred_error),
             Mask::AllFalse(_) => return Ok(self.all_null()),
             Mask::Values(valid_rows) => valid_rows,
@@ -101,13 +127,26 @@ impl RowFnExecutionArgs {
         values: ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let values = self.validate_kernel_output(values, self.row_count, ctx)?;
+        let values = self.validate_kernel_output(values, self.batch.row_count, ctx)?;
 
-        match self.validity.clone() {
+        match self.batch.validity.clone() {
             Validity::NonNullable | Validity::AllValid => {
-                self.finalize_output(values, self.row_count)
+                self.finalize_output(values, self.batch.row_count)
+            }
+            // A canonical kernel output takes the validity directly. Masking it lazily would add
+            // a scalar-function node that the caller must execute again just to attach it.
+            Validity::Array(valid) => {
+                let values = match values.as_opt::<AnyCanonical>() {
+                    Some(canonical) => mask_validity_canonical(
+                        Canonical::from(canonical),
+                        Validity::Array(valid),
+                        ctx,
+                    )?
+                    .into_array(),
+                    None => values.mask(valid)?,
+                };
+                self.finalize_output(values, self.batch.row_count)
             }
-            Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count),
             Validity::AllInvalid => {
                 unreachable!("all-invalid validity is handled before dense row execution")
             }

vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs

@@ -31,7 +31,7 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::builtins::ArrayBuiltins;
 
-impl RowFnExecutionArgs {
+impl RowFnExecutionArgs<'_> {
     /// Filter the batch to valid rows, then execute the kernel into the original row domain.
     pub(super) fn execute_filtered(
         &self,
@@ -47,16 +47,17 @@ impl RowFnExecutionArgs {
         let filter_mask = Mask::Values(MaskValuesRef::clone(valid));
 
         let filtered_inputs: SmallVec<[ArrayRef; 4]> = self
+            .batch
             .inputs
             .iter()
             .map(|input| input.filter(filter_mask.clone()))
             .collect::<VortexResult<_>>()?;
 
         let filtered_args = self.execution_args(&filtered_inputs, filtered_len);
         let values = execute_filtered_rows(filtered_args, MaskValuesRef::clone(valid), ctx)?;
-        let values = self.validate_kernel_output(values, self.row_count, ctx)?;
+        let values = self.validate_kernel_output(values, self.batch.row_count, ctx)?;
 
         let mask = valid.as_ref().into_array();
-        self.finalize_output(values.mask(mask)?, self.row_count)
+        self.finalize_output(values.mask(mask)?, self.batch.row_count)
     }
 }

vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs

@@ -3,8 +3,10 @@
 
 //! Selects a batch execution strategy.
 //!
-//! [`RowFnExecutionArgs::execute`] handles universal fast paths, then delegates to dense or
-//! valid-only execution.
+//! [`RowFnExecutionArgs::try_execute_without_retry`] handles universal fast paths and dense
+//! execution, so the visit that planned a batch can finish it without dispatching again.
+//! [`RowFnExecutionArgs::execute_partially_valid`] then delegates the remaining partially valid
+//! batches to dense retry or valid-only execution.
 
 use vortex_error::VortexResult;
 use vortex_mask::MaskValuesRef;
@@ -14,9 +16,7 @@ use super::RowPolicy;
 use super::args::BorrowedRowFnArgs;
 use crate::ArrayRef;
 use crate::ExecutionCtx;
-use crate::arrays::Constant;
 use crate::scalar_fn::unstable::row::execute::DenseAttempt;
-use crate::scalar_fn::unstable::row::types::batch_const;
 use crate::validity::Validity;
 
 mod constant;
@@ -27,16 +27,60 @@ mod valid_only;
 mod output;
 pub(crate) use output::finalize_kernel_output;
 
-impl RowFnExecutionArgs {
-    /// Apply constant folding and null handling around `kernel`.
+impl RowFnExecutionArgs<'_> {
+    /// Execute every batch that `kernel` alone can finish, or return `None` for a partially valid
+    /// batch whose plan needs a dense retry or valid-row execution.
+    ///
+    /// `kernel` reports a rejected deferred failure as a terminal error. That is correct for each
+    /// batch this method executes: no null row can be the cause of the failure.
+    pub(crate) fn try_execute_without_retry(
+        &self,
+        kernel: impl FnOnce(
+            BorrowedRowFnArgs<'_>,
+            &mut Option<Validity>,
+            &mut ExecutionCtx,
+        ) -> VortexResult<ArrayRef>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        // Strictness: an all-null batch has no observable row work. Keep the literal-constant
+        // check explicit alongside the conjoined validity invariant.
+        if matches!(self.batch.validity, Validity::AllInvalid) || self.batch.has_null_const {
+            return Ok(Some(self.all_null()));
+        }
+
+        // All inputs are constant, and their conjoined validity proves that every row is non-null.
+        // The constant check sees through extension and masked wrappers, just like argument
+        // decoding.
+        if self.batch.row_count > 0
+            && self.batch.validity.definitely_no_nulls()
+            && self.batch.all_batch_const()
+        {
+            return self.execute_all_constant(kernel, ctx).map(Some);
+        }
+
+        // Do not resolve array-backed validity for the uncommon all-valid or all-null cases here.
+        // That can execute and scan the full mask; each policy resolves it only when necessary.
+        if self.batch.validity.definitely_no_nulls() || self.plan.policy() == RowPolicy::Dense {
+            return self.execute_dense(kernel, ctx).map(Some);
+        }
+
+        Ok(None)
+    }
+
+    /// Execute a partially valid batch that
+    /// [`try_execute_without_retry`](Self::try_execute_without_retry) declined.
     ///
     /// For a partially valid batch, `try_valid_rows` can avoid filtering. `Ok(None)` filters the
     /// inputs to the valid rows and `execute_filtered_rows` runs the kernel over them while
     /// writing into the original row domain. Every result is checked against the planned shape
     /// and dtype.
-    pub(crate) fn execute(
+    pub(crate) fn execute_partially_valid(
         &self,
-        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        kernel: impl FnOnce(
+            BorrowedRowFnArgs<'_>,
+            &mut Option<Validity>,
+            &mut ExecutionCtx,
+        ) -> VortexResult<ArrayRef>,
         execute_dense_attempt: impl FnOnce(
             BorrowedRowFnArgs<'_>,
             &mut ExecutionCtx,
@@ -53,34 +97,6 @@ impl RowFnExecutionArgs {
         ) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        // Strictness: an all-null batch has no observable row work. Keep the literal-constant
-        // check explicit alongside the conjoined validity invariant.
-        if matches!(self.validity, Validity::AllInvalid)
-            || self.inputs.iter().any(|input| {
-                input
-                    .as_opt::<Constant>()
-                    .is_some_and(|constant| constant.scalar().is_null())
-            })
-        {
-            return Ok(self.all_null());
-        }
-
-        // All inputs are constant, and their conjoined validity proves that every row is non-null.
-        // The constant check sees through extension and masked wrappers, just like argument
-        // decoding.
-        if self.row_count > 0
-            && self.validity.definitely_no_nulls()
-            && self.inputs.iter().all(|input| batch_const(input).is_some())
-        {
-            return self.execute_all_constant(kernel, ctx);
-        }
-
-        // Do not resolve array-backed validity for the uncommon all-valid or all-null cases here.
-        // That can execute and scan the full mask; each policy resolves it only when necessary.
-        if self.validity.definitely_no_nulls() {
-            return self.execute_dense(kernel, ctx);
-        }
-
         match self.plan.policy() {
             RowPolicy::Dense => self.execute_dense(kernel, ctx),
             RowPolicy::DenseWithRetry => self.execute_dense_with_retry(

vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs

@@ -15,9 +15,13 @@ use crate::dtype::DType;
 use crate::scalar::Scalar;
 use crate::scalar_fn::ScalarFnId;
 
-impl RowFnExecutionArgs {
+impl RowFnExecutionArgs<'_> {
     pub(super) fn all_null(&self) -> ArrayRef {
-        ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array()
+        ConstantArray::new(
+            Scalar::null(self.result_dtype.clone()),
+            self.batch.row_count,
+        )
+        .into_array()
     }
 
     /// Label the finished column, validate it, and apply the row function's logical outer
@@ -32,7 +36,7 @@ impl RowFnExecutionArgs {
         // partially valid batch carries the same metadata as a dense one.
         let values = self.plan.relabel_output(values)?;
 
-        validate_output(self.id, &self.result_dtype, expected_len, &values)?;
+        validate_output(self.batch.id, &self.result_dtype, expected_len, &values)?;
 
         cast_output_nullability(&self.result_dtype, values)
     }
@@ -45,7 +49,7 @@ impl RowFnExecutionArgs {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
         finalize_kernel_output(
-            self.id,
+            self.batch.id,
             self.plan.storage_dtype(),
             expected_len,
             values,
@@ -68,7 +72,7 @@ pub(crate) fn finalize_kernel_output(
 ) -> VortexResult<ArrayRef> {
     validate_output(id, result_dtype, expected_len, &values)?;
     vortex_ensure!(
-        values.all_valid(ctx)?,
+        !values.dtype().is_nullable() || values.all_valid(ctx)?,
         "the {id} row kernel must produce only valid rows, got at least one null row",
     );
 

vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs

@@ -11,12 +11,17 @@ use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::builtins::ArrayBuiltins;
+use crate::validity::Validity;
 
-impl RowFnExecutionArgs {
+impl RowFnExecutionArgs<'_> {
     /// Resolve validity and try direct valid-row execution before filtered execution.
     pub(super) fn execute_valid_only(
         &self,
-        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        kernel: impl FnOnce(
+            BorrowedRowFnArgs<'_>,
+            &mut Option<Validity>,
+            &mut ExecutionCtx,
+        ) -> VortexResult<ArrayRef>,
         try_valid_rows: impl FnOnce(
             BorrowedRowFnArgs<'_>,
             MaskValuesRef,
@@ -29,15 +34,23 @@ impl RowFnExecutionArgs {
         ) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let validity = self.validity.clone().execute_mask(self.row_count, ctx)?;
+        let validity = self
+            .batch
+            .validity
+            .clone()
+            .execute_mask(self.batch.row_count, ctx)?;
 
         let valid_rows = match validity {
             // An empty mask is both all-valid and all-null. Preserve the all-valid behavior.
             Mask::AllTrue(_) | Mask::AllFalse(0) => {
-                let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?;
-                let values = self.validate_kernel_output(values, self.row_count, ctx)?;
+                let values = kernel(
+                    self.execution_args(&self.batch.inputs, self.batch.row_count),
+                    &mut None,
+                    ctx,
+                )?;
+                let values = self.validate_kernel_output(values, self.batch.row_count, ctx)?;
 
-                return self.finalize_output(values, self.row_count);
+                return self.finalize_output(values, self.batch.row_count);
             }
             Mask::AllFalse(_) => return Ok(self.all_null()),
             Mask::Values(valid_rows) => valid_rows,
@@ -62,7 +75,7 @@ impl RowFnExecutionArgs {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         let Some(values) = try_valid_rows(
-            self.execution_args(&self.inputs, self.row_count),
+            self.execution_args(&self.batch.inputs, self.batch.row_count),
             MaskValuesRef::clone(valid),
             ctx,
         )?