Add filter-and-scatter RowFn execution

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

@@ -0,0 +1,84 @@
+// 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 from the original row count to the valid row count. It runs the
+//! dense kernel on those filtered rows, then scatters the output back to the original row count.
+
+use smallvec::SmallVec;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+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::PrimitiveArray;
+use crate::dtype::Nullability;
+use crate::validity::Validity;
+
+impl RowFnExecutionArgs {
+    /// Filter the original batch to valid rows, run the dense kernel, then restore its row count.
+    pub(super) fn filter_and_scatter(
+        &self,
+        kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
+        original_validity: &Mask,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let original_len = original_validity.len();
+        let filtered_len = original_validity.true_count();
+
+        let filtered_inputs: SmallVec<[ArrayRef; 4]> = self
+            .inputs
+            .iter()
+            .map(|input| input.filter(original_validity.clone()))
+            .collect::<VortexResult<_>>()?;
+
+        let filtered_args = self.execution_args(&filtered_inputs, filtered_len);
+        let filtered = kernel(filtered_args, ctx)?;
+
+        let filtered = self.validate_kernel_output(filtered, filtered_len, ctx)?;
+
+        let output = Self::scatter_to_original_rows(filtered, original_validity)?;
+
+        self.finalize_output(output, original_len)
+    }
+
+    /// Scatter `filtered` back to the rows selected by `original_validity`.
+    ///
+    /// The result has the original length and is null at each invalid position.
+    fn scatter_to_original_rows(
+        filtered: ArrayRef,
+        original_validity: &Mask,
+    ) -> VortexResult<ArrayRef> {
+        let original_len = original_validity.len();
+
+        let AllOr::Some(valid_slices) = original_validity.slices() else {
+            // The caller handles the all-true and all-false masks.
+            vortex_bail!(
+                "filter-and-scatter requires valid and invalid rows, got an all-valid or all-invalid mask"
+            );
+        };
+
+        // Map each valid row to its position in `filtered`.
+        let mut take_indices = vec![0u64; original_len];
+
+        let valid_rows = valid_slices.iter().flat_map(|&(start, end)| start..end);
+        for (filtered_idx, original_idx) in valid_rows.enumerate() {
+            take_indices[original_idx] = u64::try_from(filtered_idx)?;
+        }
+
+        // Null indices restore invalid rows without selecting a value from `filtered`.
+        let take_indices = PrimitiveArray::new(
+            take_indices,
+            Validity::from_mask(original_validity.clone(), Nullability::Nullable),
+        )
+        .into_array();
+
+        filtered.take(take_indices)
+    }
+}

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;
@@ -29,8 +30,9 @@ 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.
+    /// For a partially valid 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,
         kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<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 and try direct valid-row execution before filter-and-scatter.
     pub(super) fn execute_valid_only(
         &self,
         kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
@@ -40,19 +34,18 @@ impl RowFnExecutionArgs {
         ) -> VortexResult<Option<ArrayRef>>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
-        let valid = match self.resolve_validity(&kernel, ctx)? {
+        let original_validity = match self.resolve_validity(&kernel, ctx)? {
             ResolvedValidity::Output(output) => return Ok(output),
-            ResolvedValidity::PartiallyValid(valid) => valid,
+            ResolvedValidity::PartiallyValid(original_validity) => original_validity,
         };
 
-        if let Some(result) = self.try_execute_valid_rows(try_valid_rows, &valid, ctx)? {
-            return Ok(result);
+        let direct_output = self.try_execute_valid_rows(try_valid_rows, &original_validity, ctx)?;
+
+        if let Some(output) = direct_output {
+            return Ok(output);
         }
 
-        vortex_panic!(
-            "valid-only execution requires direct valid-row support; {} selected an unsupported signature",
-            self.id,
-        )
+        self.filter_and_scatter(kernel, &original_validity, ctx)
     }
 
     /// Materialize validity and handle all-valid or all-null batches.
@@ -61,24 +54,24 @@ impl RowFnExecutionArgs {
         kernel: &impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ResolvedValidity> {
-        let valid = self.validity.clone().execute_mask(self.row_count, ctx)?;
+        let validity = self.validity.clone().execute_mask(self.row_count, ctx)?;
 
         // An array-backed validity can materialize to all valid even though the cheap checks in
         // `RowFnExecutionArgs::execute` could not prove that. Run the full-row kernel in that
         // case. Check all-true before all-false because an empty mask is both.
-        if valid.all_true() {
+        if validity.all_true() {
             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 = self.finalize_output(values, self.row_count)?;
 
             return Ok(ResolvedValidity::Output(values));
         }
 
-        if valid.all_false() {
+        if validity.all_false() {
             return Ok(ResolvedValidity::Output(self.all_null()));
         }
 
-        Ok(ResolvedValidity::PartiallyValid(valid))
+        Ok(ResolvedValidity::PartiallyValid(validity))
     }
 
     /// Try execution against the original inputs, then mask a returned full-length result.
@@ -89,21 +82,22 @@ impl RowFnExecutionArgs {
             &Mask,
             &mut ExecutionCtx,
         ) -> VortexResult<Option<ArrayRef>>,
-        valid: &Mask,
+        original_validity: &Mask,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
-        let Some(values) = try_valid_rows(
-            self.execution_args(&self.inputs, self.row_count),
-            valid,
-            ctx,
-        )?
-        else {
+        let row_count = original_validity.len();
+        let args = self.execution_args(&self.inputs, row_count);
+
+        let Some(values) = try_valid_rows(args, original_validity, ctx)? else {
             return Ok(None);
         };
-        let values = self.validate_kernel_output(values, valid.len(), ctx)?;
 
-        let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array();
-        self.finalize_output(values.mask(mask)?, valid.len())
-            .map(Some)
+        let values = self.validate_kernel_output(values, row_count, ctx)?;
+
+        let validity_array =
+            BoolArray::new(original_validity.to_bit_buffer(), Validity::NonNullable).into_array();
+        let masked = values.mask(validity_array)?;
+
+        self.finalize_output(masked, row_count).map(Some)
     }
 }

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

@@ -40,7 +40,7 @@ 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 for validity, constant folding, filtering, and execution.
     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

@@ -8,6 +8,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::IntoArray;
@@ -21,6 +24,7 @@ use crate::dtype::NativePType;
 use crate::dtype::Nullability;
 use crate::scalar::Scalar;
 use crate::scalar_fn::EmptyOptions;
+use crate::scalar_fn::ExecutionArgs;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::VecExecutionArgs;
 use crate::scalar_fn::unstable::row::OutputElement;
@@ -182,6 +186,34 @@ fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> 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| args.get(0),
+        |_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, filtering inputs if direct execution is unavailable.
     ValidOnly,
 }