From b6c09be0a325401e2a23f20b2ac2638e864434ea Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 20 Aug 2026 14:51:30 -0400 Subject: [PATCH 1/2] Add filter-and-scatter RowFn execution Signed-off-by: Connor Tsui --- .../row/batch/execute/filter_scatter.rs | 100 ++++++++++++++++++ .../unstable/row/batch/execute/mod.rs | 6 +- .../unstable/row/batch/execute/valid_only.rs | 54 +++++----- .../src/scalar_fn/unstable/row/batch/mod.rs | 2 +- .../src/scalar_fn/unstable/row/batch/tests.rs | 32 ++++++ .../scalar_fn/unstable/row/execute/sink.rs | 8 +- .../src/scalar_fn/unstable/row/mod.rs | 3 +- .../src/scalar_fn/unstable/row/types/sink.rs | 2 +- .../scalar_fn/unstable/row/visitor/execute.rs | 4 +- .../scalar_fn/unstable/row/visitor/plan.rs | 2 +- 10 files changed, 171 insertions(+), 42 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs new file mode 100644 index 00000000000..004773ea68f --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs @@ -0,0 +1,100 @@ +// 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::BoolArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +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, + original_validity: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + 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::>()?; + + 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 { + 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`. Invalid rows use index zero because + // their gathered values are masked below. + 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)?; + } + + let take_indices = PrimitiveArray::new(take_indices, Validity::NonNullable).into_array(); + + let expanded = filtered.take(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 expanded.dtype().is_nullable() { + let validity_array = + BoolArray::new(original_validity.to_bit_buffer(), Validity::NonNullable) + .into_array(); + + return expanded.mask(validity_array); + } + + // The gathered values are all valid, so attaching validity is sufficient. + Ok(MaskedArray::try_new( + expanded, + Validity::from_mask(original_validity.clone(), Nullability::Nullable), + )? + .into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index 60a23a1154d..28dc5ab7969 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/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, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs index c614e04dc3e..bbf9b1a7689 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs +++ b/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, @@ -40,19 +34,18 @@ impl RowFnExecutionArgs { ) -> VortexResult>, ctx: &mut ExecutionCtx, ) -> VortexResult { - 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,12 +54,12 @@ impl RowFnExecutionArgs { kernel: &impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - 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)?; @@ -74,11 +67,11 @@ impl RowFnExecutionArgs { 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>, - valid: &Mask, + original_validity: &Mask, ctx: &mut ExecutionCtx, ) -> VortexResult> { - 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) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs index dd6cfa9ba0f..4f5dee36803 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/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. diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index dfa83cc7017..ad6b25f3661 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/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(); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index 9b9131db008..238be61f8e1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/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: &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); }; diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 5551684dd4b..e55c736a1ff 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/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. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index 2bf51cf999d..747d68d47f4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -84,7 +84,7 @@ pub unsafe trait OutputSink: '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 fn(&mut Self::Rows<'a>)> { None } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index 7016aedaf23..04e88124e8d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/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; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index 75d5852d83f..7b1e230dcf1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/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, } From 685e0d28fcd455f6bb08de20c06329186b9266e6 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 21 Aug 2026 08:55:53 -0400 Subject: [PATCH 2/2] Use nullable indices to scatter row outputs Signed-off-by: Connor Tsui --- .../row/batch/execute/filter_scatter.rs | 32 +++++-------------- 1 file changed, 8 insertions(+), 24 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs index 004773ea68f..c3f542fdabd 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs @@ -17,10 +17,7 @@ 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; @@ -67,8 +64,7 @@ impl RowFnExecutionArgs { ); }; - // Map each valid row to its position in `filtered`. Invalid rows use index zero because - // their gathered values are masked below. + // 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); @@ -76,25 +72,13 @@ impl RowFnExecutionArgs { take_indices[original_idx] = u64::try_from(filtered_idx)?; } - let take_indices = PrimitiveArray::new(take_indices, Validity::NonNullable).into_array(); - - let expanded = filtered.take(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 expanded.dtype().is_nullable() { - let validity_array = - BoolArray::new(original_validity.to_bit_buffer(), Validity::NonNullable) - .into_array(); - - return expanded.mask(validity_array); - } - - // The gathered values are all valid, so attaching validity is sufficient. - Ok(MaskedArray::try_new( - expanded, + // 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()) + ) + .into_array(); + + filtered.take(take_indices) } }