Implement RowFn batch execution

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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.
+//!
+//! [`BorrowedExecutionArgs`] can point at original, sliced, or filtered arrays while retaining the
+//! dtypes, output dtype, and null policy of the original batch plan.
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::ArrayRef;
+use crate::dtype::DType;
+use crate::scalar_fn::ExecutionArgs;
+use crate::scalar_fn::unstable::row::visitor::RowPolicy;
+
+/// A borrowed [`ExecutionArgs`] view with the planning metadata selected for its row kernel.
+///
+/// `arrays` may be filtered or 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 BorrowedExecutionArgs<'a> {
+    /// The input arrays for this kernel invocation.
+    arrays: &'a [ArrayRef],
+
+    /// The number of rows in this kernel 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> BorrowedExecutionArgs<'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 BorrowedExecutionArgs<'_> {
+    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,35 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use smallvec::SmallVec;
+use vortex_error::VortexResult;
+
+use super::super::Batch;
+use super::super::args::BorrowedExecutionArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::arrays::ConstantArray;
+use crate::scalar_fn::unstable::row::execute::RowExecution;
+
+impl Batch {
+    /// Evaluate one row of constant inputs and broadcast the validated result.
+    pub(super) fn broadcast_one_row(
+        &self,
+        kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult<RowExecution>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let one_row: SmallVec<[ArrayRef; 4]> = self
+            .inputs
+            .iter()
+            .map(|input| input.slice(0..1))
+            .collect::<VortexResult<_>>()?;
+
+        let result = VortexResult::from(kernel(self.execution_args(&one_row, 1), ctx)?)?;
+        let result = self.validate_kernel_output(result, 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,54 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use super::super::Batch;
+use super::super::args::BorrowedExecutionArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::builtins::ArrayBuiltins;
+use crate::scalar_fn::unstable::row::execute::RowExecution;
+use crate::validity::Validity;
+
+impl Batch {
+    /// Run every stored payload, then attach the input validity without materializing its mask.
+    pub(super) fn execute_dense(
+        &self,
+        kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult<RowExecution>,
+        retry_deferred_error: bool,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let values = match kernel(self.execution_args(&self.inputs, self.row_count), ctx)? {
+            RowExecution::Output(values) => values,
+            RowExecution::DeferredError(error) if retry_deferred_error => {
+                let valid = self.validity.clone().execute_mask(self.row_count, ctx)?;
+
+                // Unlike `resolve_validity`, all-true preserves the deferred error and all-false
+                // suppresses evidence that came entirely from null rows. An empty loop cannot
+                // produce deferred evidence, so the ambiguous empty mask cannot reach this arm.
+                if valid.all_true() {
+                    return Err(error);
+                }
+                if valid.all_false() {
+                    return Ok(self.all_null());
+                }
+
+                // Deferred retry receives only the dense kernel. Filter first so the retry cannot
+                // evaluate null rows again.
+                return self.filter_and_scatter(kernel, &valid, ctx);
+            }
+            RowExecution::DeferredError(error) => return Err(error),
+        };
+        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 `Batch::execute`, before the kernel ran.
+            Validity::AllInvalid => Ok(self.all_null()),
+        }
+    }
+}

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

@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use smallvec::SmallVec;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure_eq;
+use vortex_mask::AllOr;
+use vortex_mask::Mask;
+
+use super::super::Batch;
+use super::super::args::BorrowedExecutionArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::arrays::BoolArray;
+use crate::arrays::MaskedArray;
+use crate::arrays::PrimitiveArray;
+use crate::builtins::ArrayBuiltins;
+use crate::dtype::Nullability;
+use crate::scalar_fn::unstable::row::execute::RowExecution;
+use crate::validity::Validity;
+
+impl Batch {
+    /// Filter to valid rows, run the kernel, then scatter into a null-padded output.
+    pub(super) fn filter_and_scatter(
+        &self,
+        kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult<RowExecution>,
+        valid: &Mask,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let filtered: SmallVec<[ArrayRef; 4]> = self
+            .inputs
+            .iter()
+            .map(|input| input.filter(valid.clone()))
+            .collect::<VortexResult<_>>()?;
+
+        let values = VortexResult::from(kernel(
+            self.execution_args(&filtered, valid.true_count()),
+            ctx,
+        )?)?;
+        let values = self.validate_kernel_output(values, valid.true_count(), ctx)?;
+
+        self.finalize_output(self.scatter_valid(values, valid)?, valid.len())
+    }
+
+    /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set
+    /// bits, producing an array of length `valid.len()` that is null at every unset position.
+    fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult<ArrayRef> {
+        vortex_ensure_eq!(
+            values.len(),
+            valid.true_count(),
+            "the {} kernel output must contain {} filtered rows, got {}",
+            self.id,
+            valid.true_count(),
+            values.len(),
+        );
+
+        let AllOr::Some(slices) = valid.slices() else {
+            // The caller handles the all-true and all-false masks.
+            vortex_bail!(
+                "scatter_valid requires valid and invalid rows, got an all-valid or all-invalid mask"
+            );
+        };
+
+        // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index
+        // 0, and any in-bounds index would do since they are masked out below (values is non-empty
+        // here).
+        let mut indices = vec![0u64; valid.len()];
+        let mut rank = 0u64;
+        for &(start, end) in slices {
+            for index in &mut indices[start..end] {
+                *index = rank;
+                rank += 1;
+            }
+        }
+        let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array();
+
+        let scattered = values.take(indices)?;
+
+        // A nullable gathered array cannot be wrapped because a `Masked` child must be all valid.
+        // The general masking pass unions its nulls with the batch validity instead.
+        if scattered.dtype().is_nullable() {
+            let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array();
+            return scattered.mask(mask);
+        }
+
+        // The gathered values are all valid, so attaching validity is sufficient.
+        Ok(MaskedArray::try_new(
+            scattered,
+            Validity::from_mask(valid.clone(), Nullability::Nullable),
+        )?
+        .into_array())
+    }
+}

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

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Selects a batch execution strategy.
+//!
+//! [`Batch::execute`] handles universal fast paths, then delegates to dense or valid-only
+//! execution.
+
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use super::Batch;
+use super::RowPolicy;
+use super::args::BorrowedExecutionArgs;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::arrays::Constant;
+use crate::scalar_fn::unstable::row::execute::RowExecution;
+use crate::scalar_fn::unstable::row::types::batch_const;
+use crate::validity::Validity;
+
+mod constant;
+mod dense;
+mod filter_scatter;
+mod valid_only;
+
+mod output;
+pub(crate) use output::finalize_kernel_output;
+
+impl Batch {
+    /// Apply constant folding and null handling around `kernel`.
+    ///
+    /// When the mask contains valid and invalid rows, `try_unfiltered` may avoid filtering.
+    /// `Ok(None)` filters the valid rows and scatters the output back. Every result is checked
+    /// against the planned shape and dtype.
+    pub(crate) fn execute(
+        &self,
+        kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult<RowExecution>,
+        try_unfiltered: impl FnOnce(
+            BorrowedExecutionArgs<'_>,
+            &Mask,
+            &mut ExecutionCtx,
+        ) -> VortexResult<Option<RowExecution>>,
+        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 constant, and their conjoined validity proves every row non-null. This sees
+        // through extension and masked wrappers just like argument decoding does.
+        if self.row_count > 0
+            && self.validity.definitely_no_nulls()
+            && self.inputs.iter().all(|input| batch_const(input).is_some())
+        {
+            return self.broadcast_one_row(kernel, ctx);
+        }
+
+        match self.policy {
+            RowPolicy::Dense => self.execute_dense(kernel, false, ctx),
+            RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx),
+            RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx),
+        }
+    }
+}