Execute spatial distance with RowFn

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

Author: connortsui20

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

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Fallback execution for row kernels that cannot operate directly on partially valid inputs.
+//!
+//! This path filters every input, runs the dense kernel, and expands its output back to the
+//! original batch shape.
+
+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::RowFnExecutionArgs;
+use super::super::args::BorrowedRowFnArgs;
+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::validity::Validity;
+
+impl RowFnExecutionArgs {
+    /// Filter to valid rows, run the kernel, then scatter into a null-padded output.
+    pub(super) fn filter_and_scatter(
+        &self,
+        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        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 = 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"
+            );
+        };
+
+        // `take` emits one row per source index. Supplying one index per original batch row expands
+        // the compact kernel output back to the batch shape. Null rows temporarily read index zero;
+        // any in-bounds index works because the result is masked below.
+        let mut source_indices = vec![0u64; valid.len()];
+        let mut rank = 0u64;
+        for &(start, end) in slices {
+            for index in &mut source_indices[start..end] {
+                *index = rank;
+                rank += 1;
+            }
+        }
+        let source_indices =
+            PrimitiveArray::new(source_indices, Validity::NonNullable).into_array();
+
+        let expanded = values.take(source_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 expanded.dtype().is_nullable() {
+            let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array();
+            return expanded.mask(mask);
+        }
+
+        // The gathered values are all valid, so attaching validity is sufficient.
+        Ok(MaskedArray::try_new(
+            expanded,
+            Validity::from_mask(valid.clone(), Nullability::Nullable),
+        )?
+        .into_array())
+    }
+}

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

@@ -20,6 +20,7 @@ use crate::validity::Validity;
 
 mod constant;
 mod dense;
+mod filter_scatter;
 mod valid_only;
 
 mod output;
@@ -30,8 +31,8 @@ impl RowFnExecutionArgs {
     /// Apply encoded reductions, constant folding, and null handling around `kernel`.
     ///
     /// `reduce` receives the original inputs before constant broadcasting. 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.
+    /// batch, `try_valid_rows` can 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,
         reduce: impl FnOnce(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<Option<ArrayRef>>,

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

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_error::VortexResult;
-use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 
 use super::super::RowFnExecutionArgs;
@@ -24,12 +23,7 @@ enum ResolvedValidity {
 }
 
 impl RowFnExecutionArgs {
-    /// Resolve validity, then execute valid rows over the original inputs.
-    ///
-    /// # Panics
-    ///
-    /// Panics if the concrete row signature cannot use direct valid-row execution. Inputs must
-    /// support null-tolerant decoding, and output sinks must initialize skipped rows.
+    /// Resolve validity, try direct valid-row execution, then fall back to filtering.
     pub(super) fn execute_valid_only(
         &self,
         kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
@@ -49,10 +43,7 @@ impl RowFnExecutionArgs {
             return Ok(result);
         }
 
-        vortex_panic!(
-            "valid-only execution requires direct valid-row support; {} selected an unsupported signature",
-            self.id,
-        )
+        self.filter_and_scatter(kernel, &valid, ctx)
     }
 
     /// Materialize validity and handle all-valid or all-null batches.

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

@@ -40,7 +40,8 @@ pub(crate) struct RowFnExecutionArgs {
     /// The number of rows in the original execution scope.
     row_count: usize,
 
-    /// The input columns, collected once for validity, constant handling, and execution.
+    /// The input columns, collected once: constant folding inspects them and the filter strategy
+    /// filters them.
     inputs: SmallVec<[ArrayRef; 4]>,
 
     /// The input dtypes, collected with the columns and reused by both planning and execution.

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

@@ -9,6 +9,9 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_session::registry::CachedId;
 
+use super::BatchPlan;
+use super::RowFnExecutionArgs;
+use super::RowPolicy;
 use super::finalize_kernel_output;
 use crate::ArrayRef;
 use crate::ExecutionCtx;
@@ -311,6 +314,35 @@ fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn test_valid_only_filters_and_scatters() -> VortexResult<()> {
+    static ID: CachedId = CachedId::new("test.filter_and_scatter");
+
+    let input = PrimitiveArray::new(
+        vec![10_i64, 20, 30, 40],
+        Validity::from_iter([true, false, true, false]),
+    )
+    .into_array();
+    let args = VecExecutionArgs::new(vec![input.clone()], 4);
+    let batch = RowFnExecutionArgs::new(*ID, &args, |_| {
+        Ok(BatchPlan {
+            output_dtype: DType::from(i64::PTYPE),
+            policy: RowPolicy::ValidOnly,
+        })
+    })?;
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = batch.execute(
+        |_args, _ctx| Ok(None),
+        |args, _ctx| Ok(args.arrays()[0].clone()),
+        |_args, _valid, _ctx| Ok(None),
+        &mut ctx,
+    )?;
+
+    assert_arrays_eq!(&actual, &input, &mut ctx);
+    Ok(())
+}
+
 #[test]
 fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> {
     let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array();

vortex-array/src/scalar_fn/unstable/row/execute/sink.rs

@@ -98,8 +98,8 @@ where
 /// Write only the rows set in `valid`, or decline when the inputs or sink cannot support
 /// skip-invalid execution.
 ///
-/// `Ok(None)` signals that direct skip-invalid execution is unavailable. Batch execution decides
-/// how to handle the decline.
+/// `Ok(None)` signals batch execution to filter every input to the valid rows, run the dense
+/// kernel, and scatter the results back into a null-padded array.
 pub(crate) fn execute_sink_valid_rows<Args, Prepared, Sink, ApplyResult, Options>(
     args: &dyn ExecutionArgs,
     valid: &Mask,
@@ -222,8 +222,8 @@ where
         return Ok(None);
     };
 
-    // Null-tolerant decoding exposes values behind nulls without filtering. Decline when any input
-    // cannot provide those values safely.
+    // Null-tolerant decoding exposes values behind nulls without filtering. Decline so batch
+    // execution can filter when any input cannot provide those values safely.
     let Some(columns) = Args::decode_null_tolerant(args, ctx)? else {
         return Ok(None);
     };

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

@@ -14,7 +14,8 @@
 //! Unlike a general strict function, a [`RowFn`] cannot produce null from valid inputs.
 //!
 //! A _partially valid_ batch contains both valid and invalid rows. _Skip-invalid_ runs the kernel
-//! only for valid rows without changing row positions.
+//! only for valid rows without changing row positions. _Filter-and-scatter_ compacts valid rows,
+//! runs the kernel, and restores their positions.
 //!
 //! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits
 //! reduce compact failure evidence without constructing errors in that loop.

vortex-array/src/scalar_fn/unstable/row/types/sink.rs

@@ -84,7 +84,7 @@ pub unsafe trait OutputSink<Options>: 'static + Sized {
     /// `Some` enables this strategy. The initializer **must** make every row safe to finish.
     /// Callbacks overwrite valid rows, and batch execution masks skipped rows.
     ///
-    /// `None` makes skip-invalid execution unavailable for this sink.
+    /// `None` makes the executor fall back to filtering the inputs.
     fn skipped_rows_initializer() -> Option<for<'a> fn(&mut Self::Rows<'a>)> {
         None
     }

vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs

@@ -4,8 +4,8 @@
 //! Visitors that execute dense and skip-invalid row loops.
 //!
 //! Each visit revalidates its concrete signature and checks that its output dtype and execution
-//! policy match the plan before entering a row loop. [`ExecuteValidRows`] can decline when the
-//! signature cannot execute over the original inputs.
+//! policy match the plan before entering a row loop. [`ExecuteValidRows`] can decline, so the batch
+//! layer filters the inputs and retries with [`ExecuteRows`].
 
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure_eq;

vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs

@@ -133,7 +133,7 @@ pub(crate) enum RowPolicy {
     /// Evaluate all rows and mask the result.
     Dense,
 
-    /// Execute only valid rows over the original inputs.
+    /// Execute only valid rows over the original inputs before filtering.
     ValidOnly,
 }