Separate RowFn batch inspection from execution planning

Inspect borrowed inputs once for lengths, dtypes, constant classification, and strict validity. Pair collected inputs with a separate execution plan while retaining the existing dispatch and retry paths.

Extracted from #9979.

Signed-off-by: Robert Kruszewski <github@robertk.io>
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: robert3005

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

@@ -10,21 +10,21 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
 
-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>,
         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), 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

@@ -14,14 +14,17 @@ use crate::builtins::ArrayBuiltins;
 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>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?;
+        let values = kernel(
+            self.execution_args(&self.batch.inputs, self.batch.row_count),
+            ctx,
+        )?;
 
         self.finalize_dense_output(values, ctx)
     }
@@ -49,8 +52,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,7 +80,7 @@ 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
@@ -101,13 +106,13 @@ 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)
             }
-            Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count),
+            Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.batch.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

@@ -14,9 +14,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,7 +25,7 @@ mod valid_only;
 mod output;
 pub(crate) use output::finalize_kernel_output;
 
-impl RowFnExecutionArgs {
+impl RowFnExecutionArgs<'_> {
     /// Apply constant folding and null handling around `kernel`.
     ///
     /// For a partially valid batch, `try_valid_rows` can avoid filtering. `Ok(None)` filters the
@@ -55,29 +53,23 @@ impl RowFnExecutionArgs {
     ) -> 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())
-            })
-        {
+        if matches!(self.batch.validity, Validity::AllInvalid) || self.batch.has_null_const {
             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())
+        if self.batch.row_count > 0
+            && self.batch.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.validity.definitely_no_nulls() {
+        if self.batch.validity.definitely_no_nulls() {
             return self.execute_dense(kernel, ctx);
         }
 

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,

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

@@ -12,7 +12,7 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::builtins::ArrayBuiltins;
 
-impl RowFnExecutionArgs {
+impl RowFnExecutionArgs<'_> {
     /// Resolve validity and try direct valid-row execution before filtered execution.
     pub(super) fn execute_valid_only(
         &self,
@@ -29,15 +29,22 @@ 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),
+                    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 +69,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,
         )?

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

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Borrowed access to the input columns of one row-function execution.
+//!
+//! [`RowInputs`] lets batch execution inspect and decode its inputs without cloning them. An
+//! internal caller that holds its operands implements it directly. [`ArrayInputs`] adapts the
+//! columns collected from a general [`ExecutionArgs`](crate::scalar_fn::ExecutionArgs).
+
+use crate::ArrayRef;
+
+/// The same-length input columns of one row-function execution, borrowed by index.
+///
+/// Batch inspection and execution share this borrowed view of the inputs.
+pub(crate) trait RowInputs {
+    /// Return the number of input columns.
+    fn num_inputs(&self) -> usize;
+
+    /// Return the number of rows in the execution scope.
+    fn row_count(&self) -> usize;
+
+    /// Return input `index`. `index` **must** be less than [`num_inputs`](Self::num_inputs).
+    fn input(&self, index: usize) -> &ArrayRef;
+}
+
+/// [`RowInputs`] over a slice of collected columns.
+pub(crate) struct ArrayInputs<'a> {
+    /// The input columns.
+    arrays: &'a [ArrayRef],
+
+    /// The number of rows in the execution scope.
+    row_count: usize,
+}
+
+impl<'a> ArrayInputs<'a> {
+    pub(crate) fn new(arrays: &'a [ArrayRef], row_count: usize) -> Self {
+        Self { arrays, row_count }
+    }
+}
+
+impl RowInputs for ArrayInputs<'_> {
+    fn num_inputs(&self) -> usize {
+        self.arrays.len()
+    }
+
+    fn row_count(&self) -> usize {
+        self.row_count
+    }
+
+    fn input(&self, index: usize) -> &ArrayRef {
+        &self.arrays[index]
+    }
+}

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

@@ -9,9 +9,10 @@
 //! propagating strict validity, selecting an execution strategy, and validating the finished
 //! output.
 //!
-//! [`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.
+//! [`RowInputs`] borrows the input columns, and [`BatchProbe`] classifies them before dispatch.
+//! [`BatchInputs`] retains the columns and their classification. [`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.
 
 use smallvec::SmallVec;
 
@@ -23,6 +24,10 @@ use crate::validity::Validity;
 mod args;
 pub(super) use args::BorrowedRowFnArgs;
 
+mod inputs;
+pub(crate) use inputs::ArrayInputs;
+pub(crate) use inputs::RowInputs;
+
 mod execute;
 pub(super) use execute::finalize_kernel_output;
 
@@ -31,8 +36,29 @@ mod planning;
 pub(super) use super::visitor::BatchPlan;
 pub(super) use super::visitor::RowPolicy;
 
-/// The same-length input columns and metadata for one row-function execution.
-pub(crate) struct RowFnExecutionArgs {
+/// What batch execution derives from its borrowed inputs before a dispatch plans them.
+///
+/// A probe holds no input column. [`BatchInputs`] retains the columns and their probe when a batch
+/// needs owned inputs for execution.
+pub(crate) struct BatchProbe {
+    /// The number of rows in the original execution scope.
+    row_count: usize,
+
+    /// The conjoined input validity. An output row is valid exactly when it is valid in every
+    /// input. Conjoining is lazy, and null handling materializes the mask only when required.
+    validity: Validity,
+
+    /// Bit `i` is set when input `i` is a batch constant, seen through extension and masked
+    /// wrappers.
+    const_mask: u32,
+
+    /// Whether any input is a null literal constant.
+    has_null_const: bool,
+}
+
+/// The same-length input columns collected for one row-function execution, with the facts of
+/// their [`BatchProbe`].
+pub(crate) struct BatchInputs {
     /// The function being executed, named in the errors this raises.
     id: ScalarFnId,
 
@@ -42,13 +68,26 @@ pub(crate) struct RowFnExecutionArgs {
     /// The input columns, collected once for validity, constant folding, filtering, and execution.
     inputs: SmallVec<[ArrayRef; 4]>,
 
-    /// The input dtypes, collected with the columns and reused by both planning and execution.
-    arg_dtypes: SmallVec<[DType; 4]>,
-
     /// The conjoined input validity. An output row is valid exactly when it is valid in every
     /// input. Conjoining is lazy, and null handling materializes the mask only when required.
     validity: Validity,
 
+    /// Bit `i` is set when input `i` is a batch constant, seen through extension and masked
+    /// wrappers.
+    const_mask: u32,
+
+    /// Whether any input is a null literal constant.
+    has_null_const: bool,
+}
+
+/// The same-length input columns and metadata for one row-function execution.
+pub(crate) struct RowFnExecutionArgs<'a> {
+    /// The collected input columns and their conjoined validity.
+    batch: &'a BatchInputs,
+
+    /// The input dtypes, collected with the columns and reused by both planning and execution.
+    arg_dtypes: &'a [DType],
+
     /// The output dtype, widened to nullable when any input is nullable. The finished column is
     /// reconciled against this dtype.
     result_dtype: DType,