Separate RowFn batch inspection from execution planning

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.probe.row_count).into_array())
     }
 }

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

@@ -18,14 +18,17 @@ 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>,
         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.probe.row_count),
+            ctx,
+        )?;
 
         self.finalize_dense_output(values, ctx)
     }
@@ -53,8 +56,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.probe.row_count),
+            ctx,
+        )?;
 
         match attempt {
             DenseAttempt::Values(values) => self.finalize_dense_output(values, ctx),
@@ -79,7 +84,12 @@ 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
+            .probe
+            .validity
+            .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
@@ -105,11 +115,11 @@ 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.probe.row_count, ctx)?;
 
-        match self.validity.clone() {
+        match self.batch.probe.validity.clone() {
             Validity::NonNullable | Validity::AllValid => {
-                self.finalize_output(values, self.row_count)
+                self.finalize_output(values, self.batch.probe.row_count)
             }
             Validity::Array(valid) => {
                 // `values.mask(valid)` builds and optimizes a ScalarFnArray to dispatch through
@@ -126,7 +136,7 @@ impl RowFnExecutionArgs {
                     .into_array(),
                     None => values.mask(valid)?,
                 };
-                self.finalize_output(values, self.row_count)
+                self.finalize_output(values, self.batch.probe.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.probe.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.probe.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,25 @@ 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.probe.validity, Validity::AllInvalid)
+            || self.batch.probe.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.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.validity.definitely_no_nulls() {
+        if self.batch.probe.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.probe.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

@@ -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,24 @@ impl RowFnExecutionArgs {
         ) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let validity = self.validity.clone().execute_mask(self.row_count, ctx)?;
+        let validity = self
+            .batch
+            .probe
+            .validity
+            .clone()
+            .execute_mask(self.batch.probe.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.probe.row_count),
+                    ctx,
+                )?;
+                let values =
+                    self.validate_kernel_output(values, self.batch.probe.row_count, ctx)?;
 
-                return self.finalize_output(values, self.row_count);
+                return self.finalize_output(values, self.batch.probe.row_count);
             }
             Mask::AllFalse(_) => return Ok(self.all_null()),
             Mask::Values(valid_rows) => valid_rows,
@@ -62,7 +71,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.probe.row_count),
             MaskValuesRef::clone(valid),
             ctx,
         )?

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

@@ -0,0 +1,71 @@
+// 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;
+use crate::dtype::DType;
+
+/// The same-length input columns of one row-function execution, borrowed by index.
+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;
+
+    /// Return the input dtypes in order, if this already holds them contiguously.
+    ///
+    /// Batch execution otherwise clones them into a collected list.
+    fn dtypes(&self) -> Option<&[DType]> {
+        None
+    }
+
+    /// 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.
+    fn as_slice(&self) -> Option<&[ArrayRef]> {
+        None
+    }
+}
+
+/// [`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]
+    }
+
+    fn as_slice(&self) -> Option<&[ArrayRef]> {
+        Some(self.arrays)
+    }
+}