Let a row dispatch declare a runtime output dtype

vortex-array/src/scalar_fn/fns/binary/numeric/row.rs

@@ -63,7 +63,7 @@ impl RowFn for NumericBinary {
         ScalarFnVTable::id(&Binary)
     }
 
-    fn dispatch<V: RowVisitor<Self::Options>>(
+    fn dispatch<V: RowVisitor>(
         &self,
         op: &Self::Options,
         args: &[DType],
@@ -89,7 +89,7 @@ fn visit_checked<T, Op, V>(visitor: V) -> VortexResult<V::VisitResult>
 where
     T: NativePType,
     Op: CheckedPrimitiveOp<T>,
-    V: RowVisitor<NumericOperator>,
+    V: RowVisitor,
 {
     visitor.visit_deferred::<(T, T), T, Op::Fail>(
         |(lhs, rhs)| Op::apply(lhs, rhs),
@@ -106,7 +106,7 @@ where
 fn visit_div<T, V>(visitor: V) -> VortexResult<V::VisitResult>
 where
     T: CheckedArithmetic,
-    V: RowVisitor<NumericOperator>,
+    V: RowVisitor,
 {
     if T::PTYPE.is_float() {
         return visit_checked::<T, CheckedDiv, V>(visitor);

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

@@ -3,22 +3,22 @@
 
 //! Execution arguments paired with the metadata selected during planning.
 //!
-//! [`BorrowedRowFnArgs`] can point at original or sliced arrays while retaining the dtypes,
-//! output dtype, and execution policy of the original batch plan.
+//! [`BorrowedRowFnArgs`] can point at original or sliced arrays while retaining the dtypes and
+//! [`BatchPlan`] of the original batch.
 
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 
-use super::RowPolicy;
+use super::BatchPlan;
 use crate::ArrayRef;
 use crate::dtype::DType;
 use crate::scalar_fn::ExecutionArgs;
 
 /// A borrowed [`ExecutionArgs`] view with the metadata selected for its row function.
 ///
-/// `arrays` can be sliced, while `dtypes` and `output_dtype` always describe the original planned
-/// batch. Keeping them together prevents an execution path from pairing an input view with
-/// unrelated planning metadata.
+/// `arrays` can be sliced, while `dtypes` and `plan` always describe the original planned batch.
+/// Keeping them together prevents an execution path from pairing an input view with unrelated
+/// planning metadata.
 #[derive(Clone, Copy)]
 pub(crate) struct BorrowedRowFnArgs<'a> {
     /// The input arrays for this row-function invocation.
@@ -30,11 +30,8 @@ pub(crate) struct BorrowedRowFnArgs<'a> {
     /// The original input dtypes used to select the row implementation.
     dtypes: &'a [DType],
 
-    /// The non-nullable dtype built by the selected output capability.
-    output_dtype: &'a DType,
-
-    /// The nullable execution policy selected during planning.
-    policy: RowPolicy,
+    /// The plan an executing dispatch must reproduce.
+    plan: &'a BatchPlan,
 }
 
 impl<'a> BorrowedRowFnArgs<'a> {
@@ -43,15 +40,13 @@ impl<'a> BorrowedRowFnArgs<'a> {
         arrays: &'a [ArrayRef],
         row_count: usize,
         dtypes: &'a [DType],
-        output_dtype: &'a DType,
-        policy: RowPolicy,
+        plan: &'a BatchPlan,
     ) -> Self {
         Self {
             arrays,
             row_count,
             dtypes,
-            output_dtype,
-            policy,
+            plan,
         }
     }
 
@@ -60,14 +55,9 @@ impl<'a> BorrowedRowFnArgs<'a> {
         self.dtypes
     }
 
-    /// Return the non-nullable dtype built by the selected output capability.
-    pub(crate) fn output_dtype(&self) -> &'a DType {
-        self.output_dtype
-    }
-
-    /// Return the nullable execution policy selected during planning.
-    pub(crate) fn policy(&self) -> RowPolicy {
-        self.policy
+    /// Return the plan an executing dispatch must reproduce.
+    pub(crate) fn plan(&self) -> &'a BatchPlan {
+        self.plan
     }
 }
 

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

@@ -73,7 +73,7 @@ impl RowFnExecutionArgs {
             return self.execute_dense(kernel, ctx);
         }
 
-        match self.policy {
+        match self.plan.policy() {
             RowPolicy::Dense => self.execute_dense(kernel, ctx),
             RowPolicy::DenseWithRetry => {
                 self.execute_dense_with_retry(execute_dense_attempt, try_valid_rows, ctx)

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

@@ -20,25 +20,37 @@ impl RowFnExecutionArgs {
         ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array()
     }
 
-    /// Validate the finished output and apply the row function's logical outer nullability.
+    /// Label the finished column, validate it, and apply the row function's logical outer
+    /// nullability.
     pub(super) fn finalize_output(
         &self,
         values: ArrayRef,
         expected_len: usize,
     ) -> VortexResult<ArrayRef> {
+        // Label before validation so the checks below see the dtype this function returns. Every
+        // batch strategy reaches this method after masking, so an empty, all-null, constant, or
+        // 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)?;
 
         cast_output_nullability(&self.result_dtype, values)
     }
 
-    /// Validate the output from a row function before batch validity is attached.
+    /// Validate the unlabelled output from a row function before batch validity is attached.
     pub(super) fn validate_kernel_output(
         &self,
         values: ArrayRef,
         expected_len: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx)
+        finalize_kernel_output(
+            self.id,
+            self.plan.storage_dtype(),
+            expected_len,
+            values,
+            ctx,
+        )
     }
 }
 

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

@@ -49,15 +49,12 @@ pub(crate) struct RowFnExecutionArgs {
     /// input. Conjoining is lazy, and null handling materializes the mask only when required.
     validity: Validity,
 
-    /// The declared output dtype, widened to nullable when any input is nullable. Kernel output is
+    /// The output dtype, widened to nullable when any input is nullable. The finished column is
     /// reconciled against this dtype.
     result_dtype: DType,
 
-    /// The non-nullable dtype the dispatched output capability builds, computed while planning.
-    output_dtype: DType,
-
-    /// How the concrete dispatch executes nullable rows.
-    policy: RowPolicy,
+    /// The storage dtype, output label, and null-handling policy selected while planning.
+    plan: BatchPlan,
 }
 
 #[cfg(test)]

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

@@ -55,8 +55,7 @@ impl RowFnExecutionArgs {
             arg_dtypes,
             validity,
             result_dtype,
-            output_dtype: plan.output_dtype,
-            policy: plan.policy,
+            plan,
         })
     }
 
@@ -66,12 +65,6 @@ impl RowFnExecutionArgs {
         arrays: &'b [ArrayRef],
         row_count: usize,
     ) -> BorrowedRowFnArgs<'b> {
-        BorrowedRowFnArgs::new(
-            arrays,
-            row_count,
-            &self.arg_dtypes,
-            &self.output_dtype,
-            self.policy,
-        )
+        BorrowedRowFnArgs::new(arrays, row_count, &self.arg_dtypes, &self.plan)
     }
 }