Plan and execute common RowFn batches in one dispatch

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

@@ -9,16 +9,21 @@ use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
+use crate::validity::Validity;
 
 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.batch.inputs, 1), ctx)?,
+            kernel(self.execution_args(&self.batch.inputs, 1), &mut None, ctx)?,
             1,
             ctx,
         )?;

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

@@ -22,14 +22,31 @@ 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> {
+        // 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 offered = matches!(self.batch.probe.validity, Validity::Array(_));
+        let mut validity = offered.then(|| self.batch.probe.validity.clone());
         let values = kernel(
             self.execution_args(&self.batch.inputs, self.batch.probe.row_count),
+            &mut validity,
             ctx,
         )?;
 
+        // A taken validity proves nothing alone. Row kernels build non-nullable storage, so only a
+        // nullable output carries the validity. Mask every other output, including one from a
+        // collector that took the validity without attaching it.
+        if offered && validity.is_none() && values.dtype().is_nullable() {
+            return self.finalize_output(values, self.batch.probe.row_count);
+        }
+
         self.finalize_dense_output(values, ctx)
     }
 
@@ -91,9 +108,9 @@ impl RowFnExecutionArgs<'_> {
             .execute_mask(self.batch.probe.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 `BatchInputs::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,

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

@@ -3,17 +3,21 @@
 
 //! Selects a batch execution strategy.
 //!
-//! [`RowFnExecutionArgs::execute`] handles universal fast paths, then delegates to dense or
-//! valid-only execution.
+//! [`BatchInputs::try_execute_without_retry`] handles universal fast paths and dense execution, so
+//! the visit that planned a batch can finish it without dispatching again. It hands every other
+//! batch back as a [`PartiallyValidBatch`], which delegates to dense retry or valid-only execution.
 
 use vortex_error::VortexResult;
 use vortex_mask::MaskValuesRef;
 
-use super::RowFnExecutionArgs;
+use super::BatchInputs;
+use super::BatchPlan;
+use super::PartiallyValidBatch;
 use super::RowPolicy;
 use super::args::BorrowedRowFnArgs;
 use crate::ArrayRef;
 use crate::ExecutionCtx;
+use crate::dtype::DType;
 use crate::scalar_fn::unstable::row::execute::DenseAttempt;
 use crate::validity::Validity;
 
@@ -25,16 +29,64 @@ mod valid_only;
 mod output;
 pub(crate) use output::finalize_kernel_output;
 
-impl RowFnExecutionArgs<'_> {
-    /// Apply constant folding and null handling around `kernel`.
+impl<'a> BatchInputs<'a> {
+    /// Execute this batch under `plan` when `kernel` alone can finish it. Otherwise return it as a
+    /// [`PartiallyValidBatch`] that 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,
+        arg_dtypes: &[DType],
+        plan: BatchPlan,
+        kernel: impl FnOnce(
+            BorrowedRowFnArgs<'_>,
+            &mut Option<Validity>,
+            &mut ExecutionCtx,
+        ) -> VortexResult<ArrayRef>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Result<ArrayRef, PartiallyValidBatch<'a>>> {
+        let args = self.with_plan(arg_dtypes, plan);
+
+        // Strictness: an all-null batch has no observable row work. Keep the literal-constant
+        // check explicit alongside the conjoined validity invariant.
+        if matches!(self.probe.validity, Validity::AllInvalid) || self.probe.has_null_const {
+            return Ok(Ok(args.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.probe.row_count > 0
+            && self.probe.validity.definitely_no_nulls()
+            && self.all_batch_const()
+        {
+            return args.execute_all_constant(kernel, ctx).map(Ok);
+        }
+
+        // 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.probe.validity.definitely_no_nulls() || args.plan.policy() == RowPolicy::Dense {
+            return args.execute_dense(kernel, ctx).map(Ok);
+        }
+
+        let plan = args.into_plan();
+
+        Ok(Err(PartiallyValidBatch { batch: self, plan }))
+    }
+}
+
+impl PartiallyValidBatch<'_> {
+    /// Execute the partially valid batch with the dense retry or valid-row strategy of its plan.
     ///
     /// 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(
-        &self,
-        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        self,
+        arg_dtypes: &[DType],
+        kernel: impl FnOnce(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
         execute_dense_attempt: impl FnOnce(
             BorrowedRowFnArgs<'_>,
             &mut ExecutionCtx,
@@ -51,40 +103,21 @@ 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.batch.probe.validity, Validity::AllInvalid)
-            || self.batch.probe.has_null_const
-        {
-            return Ok(self.all_null());
-        }
+        let Self { batch, plan } = self;
+        let args = batch.with_plan(arg_dtypes, plan);
 
-        // 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.probe.row_count > 0
-            && self.batch.probe.validity.definitely_no_nulls()
-            && self.batch.all_batch_const()
-        {
-            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.batch.probe.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(
+        match args.plan.policy() {
+            RowPolicy::Dense => {
+                unreachable!("BatchInputs::try_execute_without_retry executes every dense batch")
+            }
+            RowPolicy::DenseWithRetry => args.execute_dense_with_retry(
                 execute_dense_attempt,
                 try_valid_rows,
                 execute_filtered_rows,
                 ctx,
             ),
             RowPolicy::ValidOnly => {
-                self.execute_valid_only(kernel, try_valid_rows, execute_filtered_rows, ctx)
+                args.execute_valid_only(kernel, try_valid_rows, execute_filtered_rows, ctx)
             }
         }
     }

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

@@ -16,7 +16,7 @@ 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 ExecutionCtx) -> VortexResult<ArrayRef>,
         try_valid_rows: impl FnOnce(
             BorrowedRowFnArgs<'_>,
             MaskValuesRef,

vortex-array/src/scalar_fn/unstable/row/batch/inputs.rs

@@ -11,7 +11,10 @@ use crate::ArrayRef;
 use crate::dtype::DType;
 
 /// The same-length input columns of one row-function execution, borrowed by index.
-pub(crate) trait RowInputs {
+///
+/// The trait is `pub` only so that the sealed tuple decoding can bound on it. The private module
+/// keeps it out of the public API.
+pub trait RowInputs {
     /// Return the number of input columns.
     fn num_inputs(&self) -> usize;
 
@@ -31,7 +34,7 @@ pub(crate) trait RowInputs {
     /// Return the input columns in order, if this already holds them contiguously.
     ///
     /// The slice **must** hold exactly [`num_inputs`](Self::num_inputs) columns. Batch execution
-    /// otherwise clones the columns.
+    /// otherwise clones the columns when a visit cannot finish the batch directly.
     fn as_slice(&self) -> Option<&[ArrayRef]> {
         None
     }

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

@@ -10,14 +10,17 @@
 //! output.
 //!
 //! [`RowInputs`] borrows the input columns, and [`BatchProbe`] classifies them before dispatch.
-//! [`BatchInputs`] keeps the columns with their probe. [`BatchPlan`] carries the nullable execution
-//! strategy selected by a concrete dispatch. [`RowFnExecutionArgs`] applies that strategy, and
-//! [`BorrowedRowFnArgs`] pairs each kernel invocation with its planning metadata.
+//! [`BatchInputs`] keeps them when a visit cannot finish the batch directly. [`BatchPlan`] carries
+//! the nullable execution strategy selected by a concrete dispatch. [`RowFnExecutionArgs`] applies
+//! that strategy, and [`BorrowedRowFnArgs`] pairs each kernel invocation with its planning
+//! metadata. A [`PartiallyValidBatch`] carries a batch that needs a dense retry or valid-row
+//! execution after the planning visit.
 
 use std::borrow::Cow;
 
 use crate::ArrayRef;
 use crate::dtype::DType;
+use crate::dtype::Nullability;
 use crate::scalar_fn::ScalarFnId;
 use crate::validity::Validity;
 
@@ -38,7 +41,8 @@ pub(super) use super::visitor::RowPolicy;
 
 /// What batch execution derives from its borrowed inputs before a dispatch plans them.
 ///
-/// A probe holds no input column. [`BatchInputs`] pairs it with the columns.
+/// A probe holds no input column. A visit that can finish the batch directly decodes the borrowed
+/// inputs. Any other visit keeps them in [`BatchInputs`].
 pub(crate) struct BatchProbe {
     /// The number of rows in the original execution scope.
     row_count: usize,
@@ -53,6 +57,9 @@ pub(crate) struct BatchProbe {
 
     /// Whether any input is a null literal constant.
     has_null_const: bool,
+
+    /// The strict result nullability: nullable when any input dtype is nullable.
+    nullability: Nullability,
 }
 
 /// The same-length input columns of one row-function execution, with their [`BatchProbe`].
@@ -67,6 +74,20 @@ pub(crate) struct BatchInputs<'a> {
     probe: BatchProbe,
 }
 
+/// A partially valid batch that [`BatchInputs::try_execute_without_retry`] could not finish, with
+/// the plan that a dispatch selected for it.
+///
+/// Only that method creates this type. [`PartiallyValidBatch::execute`] therefore never receives
+/// an all-null, all-constant, or provably null-free batch, or a batch planned with
+/// [`RowPolicy::Dense`].
+pub(crate) struct PartiallyValidBatch<'a> {
+    /// The input columns and their probe.
+    batch: BatchInputs<'a>,
+
+    /// The plan that a dispatch selected for `batch`.
+    plan: BatchPlan,
+}
+
 /// The same-length input columns and metadata for one row-function execution.
 pub(crate) struct RowFnExecutionArgs<'a> {
     /// The input columns and their conjoined validity.

vortex-array/src/scalar_fn/unstable/row/batch/planning.rs

@@ -21,6 +21,7 @@ use super::args::BorrowedRowFnArgs;
 use crate::ArrayRef;
 use crate::arrays::Constant;
 use crate::dtype::DType;
+use crate::dtype::Nullability;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::unstable::row::types::is_batch_const;
 use crate::validity::Validity;
@@ -47,6 +48,7 @@ impl BatchProbe {
 
         let mut const_mask = 0u32;
         let mut has_null_const = false;
+        let mut nullability = Nullability::NonNullable;
         for index in 0..num_inputs {
             let input = inputs.input(index);
             vortex_ensure_eq!(
@@ -62,6 +64,10 @@ impl BatchProbe {
                     .as_opt::<Constant>()
                     .is_some_and(|constant| constant.scalar().is_null());
             }
+
+            if input.dtype().is_nullable() {
+                nullability = Nullability::Nullable;
+            }
         }
 
         // A null literal makes every row invalid, so skip conjoining the other validities: each
@@ -83,8 +89,39 @@ impl BatchProbe {
             validity,
             const_mask,
             has_null_const,
+            nullability,
         })
     }
+
+    /// Return the output validity of a _direct_ batch, or return the probe unchanged for any other
+    /// batch.
+    ///
+    /// A visit finishes a direct batch from its kernel output alone, once that output carries
+    /// this validity. A direct batch has no null literal, at least one input that is not
+    /// batch-constant, and at least one possibly valid row. It therefore needs no all-null or
+    /// all-constant handling. Set `dense` when the visit plans
+    /// [`RowPolicy::Dense`](super::RowPolicy::Dense). Without it, only a batch that provably has
+    /// no null row is direct, because the kernel visits every row.
+    pub(crate) fn into_direct_validity(
+        self,
+        num_inputs: usize,
+        dense: bool,
+    ) -> Result<Validity, Self> {
+        if self.has_null_const
+            || self.const_mask.count_ones() as usize == num_inputs
+            || matches!(self.validity, Validity::AllInvalid)
+            || !(dense || self.validity.definitely_no_nulls())
+        {
+            return Err(self);
+        }
+
+        Ok(self.validity.union_nullability(self.nullability))
+    }
+
+    /// Return the batch-constant classification of the inputs: bit `i` marks input `i`.
+    pub(crate) fn const_mask(&self) -> u32 {
+        self.const_mask
+    }
 }
 
 impl<'a> BatchInputs<'a> {
@@ -119,7 +156,7 @@ impl<'a> BatchInputs<'a> {
     }
 
     /// Attach the plan a dispatch selected for these inputs.
-    pub(crate) fn with_plan<'b>(
+    pub(super) fn with_plan<'b>(
         &'b self,
         arg_dtypes: &'b [DType],
         plan: BatchPlan,
@@ -136,6 +173,11 @@ impl<'a> BatchInputs<'a> {
 }
 
 impl RowFnExecutionArgs<'_> {
+    /// Return the plan, discarding the batch view.
+    pub(super) fn into_plan(self) -> BatchPlan {
+        self.plan
+    }
+
     /// Pair an input view with this batch's planning metadata.
     pub(super) fn execution_args<'b>(
         &'b self,