Implement RowFn batch execution

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

@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! 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.
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use super::RowPolicy;
+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.
+#[derive(Clone, Copy)]
+pub(crate) struct BorrowedRowFnArgs<'a> {
+    /// The input arrays for this row-function invocation.
+    arrays: &'a [ArrayRef],
+
+    /// The number of rows in this row-function invocation.
+    row_count: usize,
+
+    /// 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,
+}
+
+impl<'a> BorrowedRowFnArgs<'a> {
+    /// Pair one input view with the planning metadata selected for its batch.
+    pub(crate) fn new(
+        arrays: &'a [ArrayRef],
+        row_count: usize,
+        dtypes: &'a [DType],
+        output_dtype: &'a DType,
+        policy: RowPolicy,
+    ) -> Self {
+        Self {
+            arrays,
+            row_count,
+            dtypes,
+            output_dtype,
+            policy,
+        }
+    }
+
+    /// Return the original input dtypes used to select the row implementation.
+    pub(crate) fn dtypes(&self) -> &'a [DType] {
+        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
+    }
+}
+
+impl ExecutionArgs for BorrowedRowFnArgs<'_> {
+    fn get(&self, index: usize) -> VortexResult<ArrayRef> {
+        self.arrays.get(index).cloned().ok_or_else(|| {
+            vortex_err!(
+                "row-function input index must be less than {}, got {index}",
+                self.arrays.len(),
+            )
+        })
+    }
+
+    fn num_inputs(&self) -> usize {
+        self.arrays.len()
+    }
+
+    fn row_count(&self) -> usize {
+        self.row_count
+    }
+}

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

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use smallvec::SmallVec;
+use vortex_error::VortexResult;
+
+use super::super::RowFnExecutionArgs;
+use super::super::args::BorrowedRowFnArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::arrays::ConstantArray;
+
+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 one_row: SmallVec<[ArrayRef; 4]> = self
+            .inputs
+            .iter()
+            .map(|input| input.slice(0..1))
+            .collect::<VortexResult<_>>()?;
+
+        let result =
+            self.validate_kernel_output(kernel(self.execution_args(&one_row, 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())
+    }
+}

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

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use super::super::RowFnExecutionArgs;
+use super::super::args::BorrowedRowFnArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::builtins::ArrayBuiltins;
+use crate::validity::Validity;
+
+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 = self.validate_kernel_output(values, self.row_count, ctx)?;
+
+        match self.validity.clone() {
+            Validity::NonNullable | Validity::AllValid => {
+                self.finalize_output(values, self.row_count)
+            }
+            Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count),
+            // Handled by the guard in `RowFnExecutionArgs::execute`, before the kernel ran.
+            Validity::AllInvalid => Ok(self.all_null()),
+        }
+    }
+}

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

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Selects a batch execution strategy.
+//!
+//! [`RowFnExecutionArgs::execute`] handles universal fast paths, then delegates to dense or
+//! valid-only execution.
+
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use super::RowFnExecutionArgs;
+use super::RowPolicy;
+use super::args::BorrowedRowFnArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::arrays::Constant;
+use crate::scalar_fn::unstable::row::types::batch_const;
+use crate::validity::Validity;
+
+mod constant;
+mod dense;
+mod valid_only;
+
+mod output;
+#[cfg(test)]
+pub(crate) use output::finalize_kernel_output;
+
+impl RowFnExecutionArgs {
+    /// Apply constant folding and null handling around `kernel`.
+    ///
+    /// For a partially valid batch, `try_valid_rows` executes only valid rows over the original
+    /// inputs. Every result is checked against the planned shape and dtype.
+    pub(crate) fn execute(
+        &self,
+        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        try_valid_rows: impl FnOnce(
+            BorrowedRowFnArgs<'_>,
+            &Mask,
+            &mut ExecutionCtx,
+        ) -> VortexResult<Option<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);
+        }
+
+        // A known all-valid batch does not need to materialize validity, even when its row policy
+        // only permits valid rows.
+        if self.validity.definitely_no_nulls() {
+            return self.execute_dense(kernel, ctx);
+        }
+
+        match self.policy {
+            RowPolicy::Dense => self.execute_dense(kernel, ctx),
+            RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_valid_rows, ctx),
+        }
+    }
+}

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

@@ -0,0 +1,100 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_ensure_eq;
+
+use super::super::RowFnExecutionArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::arrays::ConstantArray;
+use crate::builtins::ArrayBuiltins;
+use crate::dtype::DType;
+use crate::scalar::Scalar;
+use crate::scalar_fn::ScalarFnId;
+
+impl RowFnExecutionArgs {
+    pub(super) fn all_null(&self) -> ArrayRef {
+        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.
+    pub(super) fn finalize_output(
+        &self,
+        values: ArrayRef,
+        expected_len: usize,
+    ) -> VortexResult<ArrayRef> {
+        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.
+    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)
+    }
+}
+
+/// Validate the output produced directly by a row function.
+///
+/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` except for
+/// outer nullability, and every produced row **must** be valid. Batch execution owns strict null
+/// propagation and attaches input-derived validity only after this boundary.
+pub(crate) fn finalize_kernel_output(
+    id: ScalarFnId,
+    result_dtype: &DType,
+    expected_len: usize,
+    values: ArrayRef,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    validate_output(id, result_dtype, expected_len, &values)?;
+    vortex_ensure!(
+        values.all_valid(ctx)?,
+        "the {id} row kernel must produce only valid rows, got at least one null row",
+    );
+
+    cast_output_nullability(result_dtype, values)
+}
+
+/// Validate an output's shape and logical dtype without executing an outer-nullability cast.
+fn validate_output(
+    id: ScalarFnId,
+    result_dtype: &DType,
+    expected_len: usize,
+    values: &ArrayRef,
+) -> VortexResult<()> {
+    vortex_ensure_eq!(
+        values.len(),
+        expected_len,
+        "the {id} kernel output must contain {expected_len} rows, got {}",
+        values.len(),
+    );
+    let values_with_result_nullability =
+        values.dtype().with_nullability(result_dtype.nullability());
+    vortex_ensure!(
+        values_with_result_nullability == *result_dtype,
+        "the {id} output dtype must match {result_dtype} except for outer nullability, got {}",
+        values.dtype(),
+    );
+
+    Ok(())
+}
+
+/// Cast only the outer output nullability after validation accepts every other dtype component.
+///
+/// This changes no logical dtype component other than outer nullability. An encoding that cannot
+/// rewrite its nullability directly can still retain a lazy cast until execution.
+fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult<ArrayRef> {
+    if values.dtype() == result_dtype {
+        Ok(values)
+    } else {
+        values.cast(result_dtype.clone())
+    }
+}