diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs index 099f0cfd135..cd5ab2aeea1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -33,7 +33,6 @@ impl RowFnExecutionArgs { /// recompute the error. pub(super) fn execute_dense_with_retry( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, execute_dense_attempt: impl FnOnce( BorrowedRowFnArgs<'_>, &mut ExecutionCtx, @@ -43,6 +42,11 @@ impl RowFnExecutionArgs { MaskValuesRef, &mut ExecutionCtx, ) -> VortexResult>, + execute_filtered_rows: impl FnOnce( + BorrowedRowFnArgs<'_>, + MaskValuesRef, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let attempt = @@ -51,7 +55,7 @@ impl RowFnExecutionArgs { match attempt { DenseAttempt::Values(values) => self.finalize_dense_output(values, ctx), DenseAttempt::DeferredError(error) => { - self.resolve_deferred_error(error, kernel, try_valid_rows, ctx) + self.resolve_deferred_error(error, try_valid_rows, execute_filtered_rows, ctx) } } } @@ -59,12 +63,16 @@ impl RowFnExecutionArgs { fn resolve_deferred_error( &self, deferred_error: VortexError, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_valid_rows: impl FnOnce( BorrowedRowFnArgs<'_>, MaskValuesRef, &mut ExecutionCtx, ) -> VortexResult>, + execute_filtered_rows: impl FnOnce( + BorrowedRowFnArgs<'_>, + MaskValuesRef, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let valid_rows = match self.validity.execute_mask(self.row_count, ctx)? { @@ -85,7 +93,7 @@ impl RowFnExecutionArgs { return Ok(result); } - self.filter_and_scatter(kernel, &valid_rows, ctx) + self.execute_filtered(execute_filtered_rows, &valid_rows, ctx) } fn finalize_dense_output( 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 deleted file mode 100644 index 988b59105e9..00000000000 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs +++ /dev/null @@ -1,95 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Fallback execution for row kernels that cannot operate on the original partially valid inputs. -//! -//! Direct skip-invalid execution retains the original row domain: it decodes every input column, -//! initializes a full-length output, and visits only valid rows. Some input representations cannot -//! decode unspecified payloads behind nulls, and some output sinks cannot initialize rows that the -//! kernel skips. In either case, there is no safe row loop to enter over the original columns. -//! -//! This fallback filters every input to the valid row domain before decoding it. The ordinary -//! kernel then produces a compact, all-valid output. A take with nullable indices scatters those -//! values back to the original row domain and restores its nulls. -//! -//! Filtering and scattering add columnar work around the row loop. Filtering nested or compressed -//! inputs can also materialize their selected representation. Batch execution therefore tries -//! direct skip-invalid execution first and uses this path only when a required capability is -//! unavailable. -//! -//! A prepared, type-specific view can provide selected-row decoding without filtering. That is an -//! optional direct-execution capability rather than a replacement for this fallback: it requires a -//! representation with safe selected access, and it does not initialize an output sink that cannot -//! represent skipped rows. Calling the general scalar-extraction API once per valid row is also not -//! equivalent because it repeats array execution and scalar construction inside the row loop. - -use smallvec::SmallVec; -use vortex_error::VortexResult; -use vortex_mask::Mask; -use vortex_mask::MaskValuesRef; - -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 kernel, then restore its row count. - pub(super) fn filter_and_scatter( - &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, - original_validity: &MaskValuesRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let original_len = original_validity.len(); - let filtered_len = original_validity.true_count(); - let filter_mask = Mask::Values(MaskValuesRef::clone(original_validity)); - - let filtered_inputs: SmallVec<[ArrayRef; 4]> = self - .inputs - .iter() - .map(|input| input.filter(filter_mask.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`. - fn scatter_to_original_rows( - filtered: ArrayRef, - original_validity: &MaskValuesRef, - ) -> VortexResult { - let original_len = original_validity.len(); - let mut take_indices = vec![0u64; original_len]; - - let valid_rows = original_validity - .slices() - .iter() - .flat_map(|&(start, end)| start..end); - for (filtered_index, original_index) in valid_rows.enumerate() { - take_indices[original_index] = u64::try_from(filtered_index)?; - } - - // Null indices restore invalid rows without selecting a value from the compact output. - let take_indices = PrimitiveArray::new( - take_indices, - Validity::from_mask( - Mask::Values(MaskValuesRef::clone(original_validity)), - Nullability::Nullable, - ), - ) - .into_array(); - - filtered.take(take_indices) - } -} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs new file mode 100644 index 00000000000..6ed3fab2a1a --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filtered.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Filtered execution for row kernels that cannot operate on the original partially valid inputs. +//! +//! Direct skip-invalid execution retains the original input columns: it decodes every column, +//! initializes a full-length output, and visits only valid rows. Some input representations cannot +//! decode unspecified payloads behind nulls, so there is no safe row loop to enter over the +//! original columns. +//! +//! This path filters every input to the valid row domain before decoding it. The kernel then +//! iterates only the valid rows: it reads consecutive rows from the compact, all-valid inputs and +//! writes each result directly at its original row index into a full-length output. Skipped +//! positions hold placeholders that batch execution masks, exactly like direct skip-invalid +//! execution, so the output never needs a columnar scatter. +//! +//! Filtering nested or compressed inputs can materialize their selected representation, so batch +//! execution tries direct skip-invalid execution first and filters only when a required input +//! capability is unavailable. Calling the general scalar-extraction API once per valid row is not +//! equivalent because it repeats array execution and scalar construction inside the row loop. + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_mask::MaskValuesRef; + +use super::super::RowFnExecutionArgs; +use super::super::args::BorrowedRowFnArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::builtins::ArrayBuiltins; + +impl RowFnExecutionArgs { + /// Filter the batch to valid rows, then execute the kernel into the original row domain. + pub(super) fn execute_filtered( + &self, + execute_filtered_rows: impl FnOnce( + BorrowedRowFnArgs<'_>, + MaskValuesRef, + &mut ExecutionCtx, + ) -> VortexResult, + valid: &MaskValuesRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered_len = valid.true_count(); + let filter_mask = Mask::Values(MaskValuesRef::clone(valid)); + + let filtered_inputs: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(filter_mask.clone())) + .collect::>()?; + + let filtered_args = self.execution_args(&filtered_inputs, filtered_len); + let values = execute_filtered_rows(filtered_args, MaskValuesRef::clone(valid), ctx)?; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + + let mask = valid.as_ref().into_array(); + self.finalize_output(values.mask(mask)?, self.row_count) + } +} 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 322cd497ac2..7298abfc9bc 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 @@ -21,7 +21,7 @@ use crate::validity::Validity; mod constant; mod dense; -mod filter_scatter; +mod filtered; mod valid_only; mod output; @@ -31,7 +31,8 @@ impl RowFnExecutionArgs { /// Apply constant folding and null handling around `kernel`. /// /// 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 + /// inputs to the valid rows and `execute_filtered_rows` runs the kernel over them while + /// writing into the original row domain. Every result is checked against the planned shape /// and dtype. pub(crate) fn execute( &self, @@ -45,6 +46,11 @@ impl RowFnExecutionArgs { MaskValuesRef, &mut ExecutionCtx, ) -> VortexResult>, + execute_filtered_rows: impl FnOnce( + BorrowedRowFnArgs<'_>, + MaskValuesRef, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { // Strictness: an all-null batch has no observable row work. Keep the literal-constant @@ -77,10 +83,15 @@ impl RowFnExecutionArgs { match self.plan.policy() { RowPolicy::Dense => self.execute_dense(kernel, ctx), - RowPolicy::DenseWithRetry => { - self.execute_dense_with_retry(kernel, execute_dense_attempt, try_valid_rows, ctx) + RowPolicy::DenseWithRetry => self.execute_dense_with_retry( + execute_dense_attempt, + try_valid_rows, + execute_filtered_rows, + ctx, + ), + RowPolicy::ValidOnly => { + self.execute_valid_only(kernel, try_valid_rows, execute_filtered_rows, ctx) } - RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_valid_rows, ctx), } } } 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 a05a476a838..62c5afb7458 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 @@ -13,7 +13,7 @@ use crate::IntoArray; use crate::builtins::ArrayBuiltins; impl RowFnExecutionArgs { - /// Resolve validity and try direct valid-row execution before filter-and-scatter. + /// Resolve validity and try direct valid-row execution before filtered execution. pub(super) fn execute_valid_only( &self, kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -22,6 +22,11 @@ impl RowFnExecutionArgs { MaskValuesRef, &mut ExecutionCtx, ) -> VortexResult>, + execute_filtered_rows: impl FnOnce( + BorrowedRowFnArgs<'_>, + MaskValuesRef, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let validity = self.validity.clone().execute_mask(self.row_count, ctx)?; @@ -42,7 +47,7 @@ impl RowFnExecutionArgs { return Ok(result); } - self.filter_and_scatter(kernel, &valid_rows, ctx) + self.execute_filtered(execute_filtered_rows, &valid_rows, ctx) } /// Try execution against the original inputs, then mask a returned full-length result. 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 8aed0dd6618..00f14e34b07 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -63,13 +63,13 @@ impl DeferredAdd { struct ValidOnlyIdentity; #[derive(Clone)] -struct FilterAndScatterIdentity; +struct FilteredIdentity; #[derive(Clone)] struct DenseRetryIncrement; #[derive(Clone)] -struct FilterAndScatterRepeat; +struct FilteredRepeat; #[derive(Clone)] struct InvalidKernelOutput; @@ -212,7 +212,8 @@ impl OutputElement for NullProducingI64 { struct I64Sink(BufferMut); // SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that -// initialized slice. The `()` write token therefore proves no additional invariant. +// initialized slice. The `()` write token therefore proves no additional invariant, and the +// default skipped-row initializer only rewrites the zeroes. unsafe impl OutputSink for I64Sink { type Params = (); type Rows<'a> = &'a mut [i64]; @@ -332,14 +333,14 @@ impl RowFn for ValidOnlyIdentity { } } -impl RowFn for FilterAndScatterIdentity { +impl RowFn for FilteredIdentity { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; const INFALLIBLE: bool = false; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("test.filter_and_scatter_identity"); + static ID: CachedId = CachedId::new("test.filtered_identity"); *ID } @@ -383,14 +384,14 @@ impl RowFn for DenseRetryIncrement { } } -impl RowFn for FilterAndScatterRepeat { +impl RowFn for FilteredRepeat { type Options = usize; const ARG_NAMES: &'static [&'static str] = &["value"]; const INFALLIBLE: bool = true; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("test.filter_and_scatter_repeat"); + static ID: CachedId = CachedId::new("test.filtered_repeat"); *ID } @@ -403,7 +404,7 @@ impl RowFn for FilterAndScatterRepeat { vortex_ensure!( u32::try_from(*width).is_ok(), InvalidArgument: - "test.filter_and_scatter_repeat width must fit in u32, got {width}", + "test.filtered_repeat width must fit in u32, got {width}", ); visitor @@ -720,14 +721,14 @@ fn test_valid_only_bool_output_skips_invalid_rows() -> VortexResult<()> { } #[test] -fn test_filter_and_scatter_skips_invalid_decode_payloads() -> VortexResult<()> { +fn test_filtered_execution_skips_invalid_decode_payloads() -> VortexResult<()> { let validity = Validity::from_iter([false, true, false, true]); let input = PrimitiveArray::new(vec![i64::MIN, 10, i64::MIN, 30], validity.clone()).into_array(); let args = VecExecutionArgs::new(vec![input], 4); let mut ctx = array_session().create_execution_ctx(); - let actual = execute_rows(&FilterAndScatterIdentity, &EmptyOptions, &args, &mut ctx)?; + let actual = execute_rows(&FilteredIdentity, &EmptyOptions, &args, &mut ctx)?; let expected = PrimitiveArray::new(vec![0_i64, 10, 0, 30], validity).into_array(); assert_arrays_eq!(&actual, &expected, &mut ctx); @@ -737,7 +738,7 @@ fn test_filter_and_scatter_skips_invalid_decode_payloads() -> VortexResult<()> { #[rstest] #[case::width_two(2, vec![0_i64, 0, 7, 7])] #[case::zero_width(0, vec![])] -fn test_filter_and_scatter_preserves_runtime_sink_params( +fn test_filtered_execution_preserves_runtime_sink_params( #[case] width: usize, #[case] expected_elements: Vec, ) -> VortexResult<()> { @@ -746,7 +747,7 @@ fn test_filter_and_scatter_preserves_runtime_sink_params( let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx(); - let actual = execute_rows(&FilterAndScatterRepeat, &width, &args, &mut ctx)?; + let actual = execute_rows(&FilteredRepeat, &width, &args, &mut ctx)?; let expected = FixedSizeListArray::new( PrimitiveArray::from_iter(expected_elements).into_array(), u32::try_from(width) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs index 3ed0b4801f4..9fba82dae43 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -8,7 +8,9 @@ mod owned; pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_filtered; pub(super) use owned::execute_owned_infallible; +pub(super) use owned::execute_owned_infallible_filtered; pub(super) use owned::execute_owned_infallible_valid_rows; pub(super) use owned::execute_owned_valid_rows; @@ -18,4 +20,5 @@ pub(super) use retry::execute_owned_dense_attempt; mod sink; pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_filtered; pub(super) use sink::execute_sink_valid_rows; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index 0224e75eb03..bab221c9b1f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -5,7 +5,9 @@ //! //! [`execute_owned`] writes fallible row results into spare vector capacity and reduces compact //! failure evidence outside the hot loop. [`execute_owned_infallible`] lets the output type map a -//! validated row source directly into its physical representation. +//! validated row source directly into its physical representation. The `_valid_rows` variants skip +//! invalid rows over the original inputs, and the `_filtered` variants read inputs filtered to the +//! valid rows while writing each output at its original row index. use std::ops::BitOrAssign; @@ -81,6 +83,106 @@ where ) } +/// Decode filtered inputs, then store one output for each valid row from an infallible kernel. +pub(crate) fn execute_owned_infallible_filtered( + args: &dyn ExecutionArgs, + valid: &MaskValuesRef, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned_filtered::( + args, + valid, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode inputs filtered to valid rows, then write one output at each valid row's original index. +/// +/// `args` addresses only the valid rows of the original batch, in order. `valid` is the original +/// batch's conjoined validity: each of its set positions receives the output of the next filtered +/// row, and unset positions keep [`Default::default`] placeholders that batch execution masks. +/// This writes directly into the original row domain, so the compact kernel output never needs a +/// columnar scatter. +pub(crate) fn execute_owned_filtered( + args: &dyn ExecutionArgs, + valid: &MaskValuesRef, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: FailureEvidence, +{ + const { assert_owned_output_needs_no_drop::() }; + + let columns = Args::decode(args, ctx)?; + + let filtered_len = args.row_count(); + vortex_ensure_eq!( + valid.true_count(), + filtered_len, + "the filtered batch must contain one row per valid row: {} valid rows, got {filtered_len}", + valid.true_count(), + ); + + let prepared = prepare(Args::const_values(&columns)); + let valid_rows = valid.bit_buffer(); + let mut values: Vec = std::iter::repeat_with(Out::default) + .take(valid_rows.len()) + .collect(); + let mut failure = Fail::default(); + let mut filtered_index = 0; + + if let Some(views) = Args::views_if_no_consts(&columns) { + vortex_ensure!( + Args::view_lens_match(&views, filtered_len), + "a decoded row input does not address exactly {filtered_len} rows", + ); + + valid_rows.for_each_set_index(|index| { + // SAFETY: the ascending set-index traversal runs exactly `true_count` times, and the + // checks above proved every view addresses `filtered_len == true_count` rows. + let elements = unsafe { Args::get_from_views_unchecked(&views, filtered_index) }; + let (value, row_failure) = apply(&prepared, elements); + + // SAFETY: every set index is below the mask length, which sized `values`. + unsafe { *values.get_unchecked_mut(index) = value }; + failure |= row_failure; + filtered_index += 1; + }); + } else { + vortex_ensure!( + Args::decoded_lens_match(&columns, filtered_len), + "a decoded row input does not address exactly {filtered_len} rows", + ); + + valid_rows.for_each_set_index(|index| { + let (value, row_failure) = apply(&prepared, Args::get(&columns, filtered_index)); + + // SAFETY: every set index is below the mask length, which sized `values`. + unsafe { *values.get_unchecked_mut(index) = value }; + failure |= row_failure; + filtered_index += 1; + }); + } + + finish_failure(failure)?; + + Ok(Out::build(values)) +} + /// Decode nullable inputs, then store outputs and combine failure evidence for valid rows. pub(crate) fn execute_owned_valid_rows( args: &dyn ExecutionArgs, 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 9b9d3b196f1..f31ab9de114 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -4,8 +4,9 @@ //! Executes row kernels that write through an [`OutputSink`]. //! //! Dense execution visits every row. Skip-invalid execution initializes skipped output rows and -//! visits only rows that are valid in every input. Skip-invalid execution declines when either the -//! input representation or sink cannot support that path. +//! visits only rows that are valid in every input. Direct skip-invalid execution declines when the +//! input representation cannot decode null payloads; filtered execution then reads inputs filtered +//! to the valid rows while still writing into the original row domain. use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -93,7 +94,7 @@ where unsafe { Sink::finish(sink) } } -/// Write only the rows set in `valid`, or decline when the inputs or sink cannot support +/// Write only the rows set in `valid`, or decline when the inputs cannot support direct /// skip-invalid execution. /// /// `Ok(None)` signals that direct skip-invalid execution is unavailable. Batch execution decides @@ -112,7 +113,6 @@ where ApplyResult: SinkResult, { let Some(ValidRowsSetup { - initialize_skipped_rows, columns, valid_rows, row_count, @@ -131,7 +131,7 @@ where { // Initialize every slot before visiting only valid rows. let mut rows = Sink::rows(&mut sink); - initialize_skipped_rows(&mut rows); + Sink::initialize_skipped_rows(&mut rows); // The initializer can change addressability. Recheck it so LLVM can prove every mask // index is in bounds. @@ -178,6 +178,99 @@ where unsafe { Sink::finish(sink) }.map(Some) } +/// Decode inputs filtered to valid rows, then write one sink row per valid row while iterating. +/// +/// `args` addresses only the valid rows of the original batch, in order. The sink covers the +/// original row domain: each set position of `valid` receives the output of the next filtered +/// row, and [`OutputSink::initialize_skipped_rows`] makes unset positions safe to finish before +/// batch execution masks them. +pub(crate) fn execute_sink_filtered( + args: &dyn ExecutionArgs, + valid: &MaskValuesRef, + params: &Sink::Params, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, +{ + let columns = Args::decode(args, ctx)?; + + let filtered_len = args.row_count(); + vortex_ensure_eq!( + valid.true_count(), + filtered_len, + "the filtered batch must contain one row per valid row: {} valid rows, got {filtered_len}", + valid.true_count(), + ); + + let original_len = valid.len(); + let mut sink = Sink::with_capacity(original_len, params)?; + + let valid_rows = valid.bit_buffer(); + let views = Args::views_if_no_consts(&columns); + let const_values = Args::const_values(&columns); + let prepared = prepare(const_values); + + // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. + { + // Initialize every slot before visiting only valid rows. + let mut rows = Sink::rows(&mut sink); + Sink::initialize_skipped_rows(&mut rows); + + // The initializer can change addressability. Recheck it so LLVM can prove every mask + // index is in bounds. + let initialized_row_count = rows.len(); + vortex_ensure_eq!( + initialized_row_count, + original_len, + "the initialized output sink must address exactly {original_len} rows, got {initialized_row_count}", + ); + + let mut filtered_index = 0; + if let Some(views) = views { + if !Args::view_lens_match(&views, filtered_len) { + decoded_length_error(filtered_len)?; + } + + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check proved that the sink addresses + // every mask index, which is below the mask's length. + let output = unsafe { Sink::row_unchecked(&mut rows, index) }; + + // SAFETY: the ascending set-index traversal runs at most `true_count` times, and + // the checks above proved every view addresses `filtered_len == true_count` rows. + let elements = unsafe { Args::get_from_views_unchecked(&views, filtered_index) }; + filtered_index += 1; + + apply(&prepared, elements, output).into_result() + })?; + } else { + if !Args::decoded_lens_match(&columns, filtered_len) { + decoded_length_error(filtered_len)?; + } + + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check proved that the sink addresses + // every mask index, which is below the mask's length. + let output = unsafe { Sink::row_unchecked(&mut rows, index) }; + + let elements = Args::get(&columns, filtered_index); + filtered_index += 1; + + apply(&prepared, elements, output).into_result() + })?; + } + } + + // SAFETY: the initializer completed before traversal, and every visited callback completed + // successfully and returned the required write token. + unsafe { Sink::finish(sink) } +} + /// Construct a decoded-length error outside the traversal branches. /// /// Owned execution (`owned.rs`) derives its index from an output-slice iterator. Sink execution @@ -196,14 +289,13 @@ where Args: ElementTuple, Sink: OutputSink, { - initialize_skipped_rows: for<'rows> fn(&mut Sink::Rows<'rows>), columns: Args::Columns, valid_rows: &'valid BitBuffer, row_count: usize, sink: Sink, } -/// Resolve the capabilities, inputs, sink, and validity mask for skip-invalid execution. +/// Resolve the inputs, sink, and validity mask for direct skip-invalid execution. fn setup_sink_valid_rows<'valid, Args, Sink>( args: &dyn ExecutionArgs, valid: &'valid MaskValuesRef, @@ -214,11 +306,6 @@ where Args: ElementTuple, Sink: OutputSink, { - // The initializer both declares support for skipping rows and initializes those rows. - let Some(initialize_skipped_rows) = Sink::skipped_rows_initializer() else { - return Ok(None); - }; - // Null-tolerant decoding exposes values behind nulls without filtering. Decline when any input // cannot provide those values safely. let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { @@ -241,7 +328,6 @@ where ); Ok(Some(ValidRowsSetup { - initialize_skipped_rows, columns, valid_rows, row_count, @@ -253,7 +339,6 @@ where mod tests { use vortex_error::VortexResult; use vortex_error::vortex_bail; - use vortex_error::vortex_err; use vortex_mask::Mask; use super::execute_sink_valid_rows; @@ -266,40 +351,9 @@ mod tests { use crate::dtype::NativePType; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::OutputSink; - use crate::validity::Validity; - - struct NonSkippingSink; struct ShrinkingSink(Vec); - // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or - // `finish` through the executor. The row-initialization requirements are therefore vacuous. - unsafe impl OutputSink for NonSkippingSink { - type Params = (); - type Rows<'a> = (); - type Row<'a> = (); - type WriteToken = (); - - fn storage_dtype(_params: &Self::Params) -> DType { - DType::from(i64::PTYPE) - } - - fn with_capacity(_rows: usize, _params: &Self::Params) -> VortexResult { - Err(vortex_err!( - "a non-skipping sink must decline before allocation" - )) - } - - fn rows(&mut self) -> Self::Rows<'_> {} - - unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { - } - - unsafe fn finish(self) -> VortexResult { - Err(vortex_err!("a non-skipping sink must not finish")) - } - } - // SAFETY: the initializer deliberately shrinks the row collection to exercise the executor's // post-initialization length check. If execution incorrectly continues, safe indexing in // `row_unchecked` panics instead of accessing invalid memory. @@ -309,10 +363,8 @@ mod tests { type Row<'a> = &'a mut i64; type WriteToken = (); - fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { - Some(|rows| { - rows.pop(); - }) + fn initialize_skipped_rows(rows: &mut Self::Rows<'_>) { + rows.pop(); } fn storage_dtype(_params: &Self::Params) -> DType { @@ -336,29 +388,6 @@ mod tests { } } - #[test] - fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { - let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); - let args = VecExecutionArgs::new(vec![input], 2); - let Mask::Values(valid) = Mask::from_iter([true, false]) else { - vortex_bail!("the test validity must be partially valid"); - }; - let mut ctx = array_session().create_execution_ctx(); - - let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>( - &args, - &valid, - &(), - &mut ctx, - |_| (), - |_, _, _| (), - )?; - - assert!(execution.is_none()); - - Ok(()) - } - #[test] fn test_skip_invalid_sink_rechecks_rows_after_initialization() -> VortexResult<()> { let input = PrimitiveArray::from_iter([10_i64, 20]).into_array(); diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 8bf45929163..ee7d2014374 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -31,6 +31,7 @@ pub use row_fn::RowFn; mod types; pub use types::ElementTuple; pub use types::FailureEvidence; +pub use types::FillDefault; pub use types::FixedSizeListSink; pub use types::IndexedElementTuple; pub use types::InitializedElement; @@ -38,6 +39,7 @@ pub use types::InitializedRow; pub use types::InputElement; pub use types::OutputElement; pub use types::OutputSink; +pub use types::Preinitialized; pub use types::SinkResult; pub use types::UninitElementSink; pub use types::ViewLen; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/fill.rs b/vortex-array/src/scalar_fn/unstable/row/types/fill.rs new file mode 100644 index 00000000000..40bd483aab1 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/fill.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Default filling for row-loop output storage. +//! +//! [`FillDefault`] lets the default [`OutputSink::initialize_skipped_rows`] zero-initialize a +//! sink's rows without knowing their representation. [`Preinitialized`] satisfies the same bound +//! with a no-op for storage that is fully initialized at construction. +//! +//! [`OutputSink::initialize_skipped_rows`]: crate::scalar_fn::unstable::row::OutputSink::initialize_skipped_rows + +use super::ViewLen; + +/// Output row storage that can fill itself with default placeholder values. +/// +/// [`OutputSink::Rows`] requires this so the default +/// [`OutputSink::initialize_skipped_rows`] can make skipped rows safe to finish. The written +/// values are placeholders only: valid rows are overwritten by the kernel, and batch execution +/// masks skipped rows before the output is observable. +/// +/// The blanket slice implementation writes `T::default()` into every element. Coherence with that +/// blanket prevents an implementation for raw `MaybeUninit` slices, so a sink exposing +/// uninitialized storage implements this on its own row view type, like +/// [`UninitElementSink`](crate::scalar_fn::unstable::row::UninitElementSink). +/// +/// [`OutputSink::Rows`]: crate::scalar_fn::unstable::row::OutputSink::Rows +/// [`OutputSink::initialize_skipped_rows`]: crate::scalar_fn::unstable::row::OutputSink::initialize_skipped_rows +pub trait FillDefault { + /// Fill every element with its default value. + fn fill_default(&mut self); +} + +impl FillDefault for () { + fn fill_default(&mut self) {} +} + +impl FillDefault for &mut T { + fn fill_default(&mut self) { + T::fill_default(self) + } +} + +impl FillDefault for [T] { + fn fill_default(&mut self) { + self.fill(T::default()); + } +} + +impl FillDefault for Vec { + fn fill_default(&mut self) { + self.as_mut_slice().fill_default(); + } +} + +/// Row storage whose construction already initialized every row. +/// +/// A sink whose `with_capacity` returns fully initialized rows can wrap its row view in this type. +/// It satisfies the [`FillDefault`] bound on [`OutputSink::Rows`] with a no-op, so the default +/// [`OutputSink::initialize_skipped_rows`] skips redundant filling, and element types without a +/// [`Default`] implementation need none. +/// +/// [`OutputSink::Rows`]: crate::scalar_fn::unstable::row::OutputSink::Rows +/// [`OutputSink::initialize_skipped_rows`]: crate::scalar_fn::unstable::row::OutputSink::initialize_skipped_rows +pub struct Preinitialized(pub R); + +impl ViewLen for Preinitialized { + fn len(&self) -> usize { + self.0.len() + } +} + +impl FillDefault for Preinitialized { + fn fill_default(&mut self) {} +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index 19a55500e18..343b4df8bf7 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -3,10 +3,14 @@ //! Input decoding and output construction for row functions. //! -//! [`ViewLen`] reports the rows addressable through a row-loop view. [`element`] defines the Rust -//! values decoded from input columns and built into simple output columns. [`sink`] handles outputs -//! that need row handles or batch-wide state. [`result`] defines immediate and deferred row -//! outcomes. +//! [`ViewLen`] reports the rows addressable through a row-loop view, and [`FillDefault`] fills a +//! view with placeholder values. [`element`] defines the Rust values decoded from input columns +//! and built into simple output columns. [`sink`] handles outputs that need row handles or +//! batch-wide state. [`result`] defines immediate and deferred row outcomes. + +mod fill; +pub use fill::FillDefault; +pub use fill::Preinitialized; mod element; pub use element::ElementTuple; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs index 6754983e2f8..25f9d417e78 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs @@ -19,6 +19,7 @@ use crate::IntoArray; use crate::arrays::FixedSizeListArray; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::FillDefault; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::ViewLen; use crate::validity::Validity; @@ -63,6 +64,14 @@ impl ViewLen for FixedSizeRows<'_, T> { } } +impl FillDefault for FixedSizeRows<'_, T> { + fn fill_default(&mut self) { + for element in self.elements.iter_mut() { + element.write(T::default()); + } + } +} + /// A fixed-size-list sink whose row width is supplied at dispatch time. /// /// The row closure must return the [`InitializedRow`] from [`InitializedRow::fill`] on success. @@ -83,23 +92,15 @@ pub struct FixedSizeListSink { // SAFETY: `with_capacity` reserves `row_count * width` elements, and `FixedSizeRows` retains that // shape for its lifetime. Each row is one disjoint `width`-element slice. `InitializedRow::fill` -// writes every element before returning its private token, and the skipped-row initializer writes -// every flat element before masked traversal. `values` retains length zero until every row is safe -// to publish in `finish`. +// writes every element before returning its private token, and `FixedSizeRows::fill_default` +// writes every flat element before masked traversal. `values` retains length zero until every row +// is safe to publish in `finish`. unsafe impl OutputSink for FixedSizeListSink { type Params = usize; type Rows<'a> = FixedSizeRows<'a, T>; type Row<'a> = &'a mut [MaybeUninit]; type WriteToken = InitializedRow; - fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { - Some(|rows| { - for element in rows.elements.iter_mut() { - element.write(T::default()); - } - }) - } - fn storage_dtype(params: &Self::Params) -> DType { DType::FixedSizeList( Arc::new(T::element_dtype()), diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs index df13ca18e88..214defaa4e3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs @@ -10,6 +10,7 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::dtype::DType; +use crate::scalar_fn::unstable::row::FillDefault; use crate::scalar_fn::unstable::row::ViewLen; mod fixed_size_list; @@ -31,7 +32,7 @@ pub use uninit_element::UninitElementSink; /// labels the column this sink builds. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. -/// Execution can omit invalid rows when [`skipped_rows_initializer`] returns an initializer. +/// Skip-invalid execution runs [`initialize_skipped_rows`] first and then visits only valid rows. /// /// # Errors /// @@ -48,7 +49,7 @@ pub use uninit_element::UninitElementSink; /// - A borrowed [`Rows`] view **must** retain its length and index-to-row mapping until it is /// dropped. Calls to [`row_unchecked`](Self::row_unchecked) and safe uses of a returned /// [`Row`](Self::Row) **must** preserve both properties. -/// - [`skipped_rows_initializer`] is the only exception to this stability requirement. The executor +/// - [`initialize_skipped_rows`] is the only exception to this stability requirement. The executor /// checks the length again after the initializer. The initializer **must** initialize every row. /// - A row must either be initialized before the callback or require a /// [`WriteToken`] that safe code cannot produce without initializing that exact row. Evidence for @@ -56,8 +57,8 @@ pub use uninit_element::UninitElementSink; /// - `Self` and every borrowed [`Rows`] view **must** remain safe to drop if decoding, /// preparation, skipped-row initialization, or a row callback returns an error or unwinds. The /// executor can abandon a sink after any prefix of rows. -/// - [`finish`] **must** be sound once every visited callback returned its required token and the -/// skipped-row initializer, when present, ran successfully. +/// - [`finish`] **must** be sound once every visited callback returned its required token and, for +/// a skip-invalid traversal, [`initialize_skipped_rows`] ran successfully first. /// - Violating these requirements can cause undefined behavior. /// /// [`Rows`]: Self::Rows @@ -66,7 +67,7 @@ pub use uninit_element::UninitElementSink; /// [`RowFn::INFALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::INFALLIBLE /// [`RowVisitor::with_output_dtype`]: crate::scalar_fn::unstable::row::RowVisitor::with_output_dtype /// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult -/// [`skipped_rows_initializer`]: Self::skipped_rows_initializer +/// [`initialize_skipped_rows`]: Self::initialize_skipped_rows pub unsafe trait OutputSink: 'static + Sized { /// Physical parameters required to construct this sink before the row loop. /// @@ -77,8 +78,11 @@ pub unsafe trait OutputSink: 'static + Sized { /// A loop-local view of all output rows. /// /// Borrowed once before execution so the sink's buffer descriptor and shape become loop - /// invariants rather than being re-read through `&mut Self` for every row. - type Rows<'a>: ViewLen + /// invariants rather than being re-read through `&mut Self` for every row. The [`FillDefault`] + /// bound lets the default [`initialize_skipped_rows`](Self::initialize_skipped_rows) + /// zero-initialize the rows; storage that is fully initialized at construction can wrap itself + /// in [`Preinitialized`](super::Preinitialized) to make that a no-op. + type Rows<'a>: ViewLen + FillDefault where Self: 'a; @@ -94,16 +98,23 @@ pub unsafe trait OutputSink: 'static + Sized { /// **must not** be able to construct one without establishing the invariant. type WriteToken: 'static; - /// The operation that initializes every output position before + /// Initialize every output position before /// [skip-invalid execution](crate::scalar_fn::unstable::row). /// - /// `Some` enables this strategy. The initializer **must** make every row safe to finish. - /// Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// This **must** make every row safe to finish. The values are placeholders only: callbacks + /// overwrite valid rows, and batch execution masks skipped rows before the output is + /// observable, so any well-formed value works. Sink storage is non-nullable, so nulls do not + /// exist at this level. /// - /// `None` makes direct skip-invalid execution unavailable. Batch execution can instead filter - /// its inputs and run the ordinary dense sink over only valid rows. - fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { - None + /// The default implementation zero-initializes the rows through [`FillDefault`], which + /// [`Rows`](Self::Rows) provides: a plain slice of `Default` elements fills itself, a custom + /// row view implements the filling for its own storage, and rows that are fully initialized at + /// construction wrap themselves in [`Preinitialized`](super::Preinitialized) to skip it. + /// Override this method only when one batch-wide pass over the rows is the wrong operation. + /// + /// Dense execution never calls this method because it visits every row. + fn initialize_skipped_rows(rows: &mut Self::Rows<'_>) { + rows.fill_default(); } /// The dtype of the column this sink builds. @@ -132,8 +143,7 @@ pub unsafe trait OutputSink: 'static + Sized { /// # Safety /// /// The executor must have completed every row callback successfully, and each callback must - /// have returned this sink's [`WriteToken`](Self::WriteToken). When skipped rows are allowed, - /// the initializer returned by - /// [`skipped_rows_initializer`](Self::skipped_rows_initializer) must have run before traversal. + /// have returned this sink's [`WriteToken`](Self::WriteToken). When rows are skipped, + /// [`initialize_skipped_rows`](Self::initialize_skipped_rows) must have run before traversal. unsafe fn finish(self) -> VortexResult; } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs index 4087ee94fae..984d133d991 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs @@ -13,7 +13,9 @@ use vortex_error::VortexResult; use super::OutputSink; use crate::ArrayRef; use crate::dtype::DType; +use crate::scalar_fn::unstable::row::FillDefault; use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::ViewLen; /// Proof that one uninitialized element row was initialized. /// @@ -62,24 +64,38 @@ pub struct UninitElementSink { row_count: usize, } +/// A loop-local view of the uninitialized row slots of an [`UninitElementSink`]. +/// +/// This wrapper carries the [`FillDefault`] behavior a raw `MaybeUninit` slice cannot: coherence +/// with the initialized-slice blanket rules that implementation out, and filling with +/// `MaybeUninit::uninit()` would be wrong anyway. Filling writes `T::default()` into every slot. +pub struct UninitElementRows<'a, T>(&'a mut [MaybeUninit]); + +impl ViewLen for UninitElementRows<'_, T> { + fn len(&self) -> usize { + self.0.len() + } +} + +impl FillDefault for UninitElementRows<'_, T> { + fn fill_default(&mut self) { + for slot in self.0.iter_mut() { + slot.write(T::default()); + } + } +} + // SAFETY: the row slice covers exactly the reserved spare-capacity range, so each accepted index // names one distinct slot. Safe code cannot construct `InitializedElement`. Its unsafe constructor -// writes the supplied slot and requires the caller to return that exact evidence. The -// skipped-row initializer writes `T::default()` into every slot before masked traversal. +// writes the supplied slot and requires the caller to return that exact evidence. The default +// skipped-row initializer fills every slot with `T::default()` through +// `UninitElementRows::fill_default` before masked traversal. unsafe impl OutputSink for UninitElementSink { type Params = (); - type Rows<'a> = &'a mut [MaybeUninit]; + type Rows<'a> = UninitElementRows<'a, T>; type Row<'a> = &'a mut MaybeUninit; type WriteToken = InitializedElement; - fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { - Some(|rows| { - for row in rows.iter_mut() { - row.write(T::default()); - } - }) - } - fn storage_dtype(_params: &Self::Params) -> DType { T::element_dtype() } @@ -92,12 +108,12 @@ unsafe impl OutputSink for UninitElementSink< } fn rows(&mut self) -> Self::Rows<'_> { - &mut self.values.spare_capacity_mut()[..self.row_count] + UninitElementRows(&mut self.values.spare_capacity_mut()[..self.row_count]) } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { // SAFETY: required by this method's contract. - unsafe { rows.get_unchecked_mut(index) } + unsafe { rows.0.get_unchecked_mut(index) } } unsafe fn finish(mut self) -> VortexResult { 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 0501d6fc6a0..8d319c802ac 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -5,7 +5,8 @@ //! //! Each visit revalidates its concrete signature and checks that its plan matches the one planning //! selected before entering a row loop. [`ExecuteValidRows`] can decline when the signature cannot -//! execute over the original inputs. +//! execute over the original inputs; [`ExecuteFilteredRows`] then executes over inputs filtered to +//! the valid rows while writing into the original row domain. use std::marker::PhantomData; @@ -33,10 +34,13 @@ use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::SinkResult; use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_filtered; use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible_filtered; use crate::scalar_fn::unstable::row::execute::execute_owned_infallible_valid_rows; use crate::scalar_fn::unstable::row::execute::execute_owned_valid_rows; use crate::scalar_fn::unstable::row::execute::execute_sink; +use crate::scalar_fn::unstable::row::execute::execute_sink_filtered; use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; /// The runtime visit that decodes every column once and runs the selected row loop. @@ -302,3 +306,147 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { ) } } + +/// The runtime visit that executes valid rows over inputs filtered to the valid row domain. +/// +/// Batch execution filters the inputs when a representation cannot decode null payloads. This +/// visit reads consecutive filtered rows and writes each output at its original row index, so the +/// result never needs a columnar scatter. Owned outputs initialize skipped positions with +/// [`Default::default`]. Output sinks use their own skipped-row initializer. +pub(crate) struct ExecuteFilteredRows<'args, 'ctx, F: RowFn> { + /// The inputs filtered to the valid rows of the original batch. + args: &'args dyn ExecutionArgs, + + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The plan selected by the planning visit, which this visit must reproduce. + plan: &'args BatchPlan, + + /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. + output_dtype: Option, + + /// The conjoined validity over the original row domain. + valid: MaskValuesRef, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// Ties this visit to the function used by its compile-time contract checks. + function: PhantomData, +} + +impl<'args, 'ctx, F: RowFn> ExecuteFilteredRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + plan: &'args BatchPlan, + valid: MaskValuesRef, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + plan, + output_dtype: None, + valid, + ctx, + function: PhantomData, + } + } +} + +impl private::Sealed for ExecuteFilteredRows<'_, '_, F> {} + +impl RowVisitor for ExecuteFilteredRows<'_, '_, F> { + type VisitResult = ArrayRef; + + fn with_output_dtype(mut self, dtype: DType) -> Self { + self.output_dtype = Some(dtype); + self + } + + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + let visited = BatchPlan::new( + validate_owned_visit::(self.dtypes)?, + self.output_dtype, + RowPolicy::for_owned_output::(), + )?; + self.plan.ensure_reproduced_by(&visited)?; + + execute_owned_infallible_filtered::( + self.args, + &self.valid, + self.ctx, + prepare, + apply, + ) + } + + fn visit_prepared_into( + self, + params: Sink::Params, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + let visited = BatchPlan::new( + validate_sink_visit::(self.dtypes, ¶ms)?, + self.output_dtype, + RowPolicy::for_sink::(), + )?; + self.plan.ensure_reproduced_by(&visited)?; + + execute_sink_filtered::( + self.args, + &self.valid, + ¶ms, + self.ctx, + prepare, + apply, + ) + } + + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: FailureEvidence, + { + const { assert_deferred_visit_contract::() }; + let visited = BatchPlan::new( + validate_owned_visit::(self.dtypes)?, + self.output_dtype, + RowPolicy::for_deferred_output::(), + )?; + self.plan.ensure_reproduced_by(&visited)?; + + execute_owned_filtered::( + self.args, + &self.valid, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index 2c90d8a5a08..dc8d0209374 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -9,6 +9,7 @@ mod check; pub(super) use check::assert_owned_output_needs_no_drop; mod execute; +pub(super) use execute::ExecuteFilteredRows; pub(super) use execute::ExecuteRows; pub(super) use execute::ExecuteValidRows; diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 2b483d6109b..028bf536dde 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -19,6 +19,7 @@ use super::batch::finalize_kernel_output; use super::row_fn::RowFn; use super::visitor::BatchPlanner; use super::visitor::ExecuteDenseWithRetry; +use super::visitor::ExecuteFilteredRows; use super::visitor::ExecuteRows; use super::visitor::ExecuteValidRows; use crate::ArrayRef; @@ -126,6 +127,7 @@ pub fn execute_rows( |args, ctx| execute_row_kernel(function, options, args, ctx), |args, ctx| execute_dense_attempt(function, options, args, ctx), |args, valid, ctx| try_execute_valid_rows(function, options, args, valid, ctx), + |args, valid, ctx| execute_filtered_rows(function, options, args, valid, ctx), ctx, ) } @@ -200,6 +202,21 @@ fn try_execute_valid_rows( ) } +/// Execute valid rows over `args` filtered to the valid row domain of `valid`. +fn execute_filtered_rows( + function: &F, + options: &F::Options, + args: BorrowedRowFnArgs<'_>, + valid: MaskValuesRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + function.dispatch( + options, + args.dtypes(), + ExecuteFilteredRows::::new(&args, args.dtypes(), args.plan(), valid, ctx), + ) +} + fn prepare_batch( function: &F, options: &F::Options, diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index 1f315e7accb..04684c83d4e 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -13,6 +13,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::unstable::row::InputElement; use vortex_array::scalar_fn::unstable::row::OutputSink; +use vortex_array::scalar_fn::unstable::row::Preinitialized; use vortex_array::scalar_fn::unstable::row::RowVisitor; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -123,18 +124,15 @@ fn empty_polygon() -> GeoPolygon { } // SAFETY: `with_capacity` creates one initialized polygon per output row, and `Rows` is the -// corresponding mutable slice. Every in-bounds index therefore names one distinct initialized -// polygon. The sink remains safe to finish or drop after any row prefix. +// corresponding mutable slice, marked `Preinitialized` so the default skipped-row initializer +// leaves it untouched. Every in-bounds index therefore names one distinct initialized polygon. +// The sink remains safe to finish or drop after any row prefix. unsafe impl OutputSink for PolygonSink { type Params = (); - type Rows<'a> = &'a mut [GeoPolygon]; + type Rows<'a> = Preinitialized<&'a mut [GeoPolygon]>; type Row<'a> = &'a mut GeoPolygon; type WriteToken = (); - fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { - Some(|_| {}) - } - fn storage_dtype((): &Self::Params) -> DType { polygon_storage_dtype(Dimension::Xy, Nullability::NonNullable) } @@ -146,12 +144,12 @@ unsafe impl OutputSink for PolygonSink { } fn rows(&mut self) -> Self::Rows<'_> { - self.polygons.as_mut_slice() + Preinitialized(self.polygons.as_mut_slice()) } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { // SAFETY: required by this method's contract. - unsafe { rows.get_unchecked_mut(index) } + unsafe { rows.0.get_unchecked_mut(index) } } unsafe fn finish(self) -> VortexResult {