diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs new file mode 100644 index 00000000000..73722549f41 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-loop execution for owned outputs and output sinks. +//! +//! [`owned`] stores one independent value per row and reduces compact failure evidence. [`sink`] +//! drives output builders whose row handles may share batch state. Both return [`RowExecution`], +//! which distinguishes a completed array from a deferred error that batch validity may suppress. + +mod owned; +pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_infallible; + +mod outcome; +pub use outcome::RowExecution; + +mod sink; +pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_valid_rows; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs new file mode 100644 index 00000000000..fc013e7a317 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The result of a completed row loop before batch-level null handling. +//! +//! [`RowExecution`] preserves deferred failure evidence until batch execution can determine whether +//! the failing payload belonged to a valid row. + +use vortex_error::VortexError; +use vortex_error::VortexResult; + +use crate::ArrayRef; + +/// The outcome of a row loop before batch execution decides whether an error is observable. +/// +/// Together with the surrounding [`VortexResult`], this represents three outcomes: +/// +/// - `Err(error)` is a non-retryable execution or immediate row error. +/// - [`Output`](Self::Output) is a successful row loop. +/// - [`DeferredError`](Self::DeferredError) is failure evidence from a completed row loop. +/// +/// A dense loop can evaluate null payloads, so its deferred error is not always observable. Batch +/// execution can retry only valid rows to discard errors caused by null payloads. A plain +/// `VortexResult` cannot distinguish these errors from failures that a retry cannot fix. +/// +/// Once execution is known to contain only valid rows, converting this outcome into a +/// `VortexResult` turns [`DeferredError`](Self::DeferredError) into an ordinary error. +pub enum RowExecution { + /// The successfully built, full-length output column. + Output(ArrayRef), + + /// An error constructed from failure evidence reduced across a completed row loop. + DeferredError(VortexError), +} + +impl From for VortexResult { + fn from(execution: RowExecution) -> Self { + match execution { + RowExecution::Output(output) => Ok(output), + RowExecution::DeferredError(error) => Err(error), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs new file mode 100644 index 00000000000..513178cd97c --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Executes row kernels that return one independent owned value per row. +//! +//! [`execute_owned`] decodes inputs once, prepares constant state, writes into spare vector +//! capacity, and reduces compact failure evidence without putting error construction in the hot +//! loop. [`execute_owned_infallible`] removes that failure path for infallible kernels. + +use std::ops::BitOrAssign; + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::FailureEvidence; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::ViewLen; +use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; + +/// Zero-sized failure accumulator for infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column, then store one output per row from an infallible kernel. +pub(crate) fn execute_owned_infallible( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned::( + args, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode every input column, then store outputs and combine per-row failure evidence. +pub(crate) fn execute_owned( + args: &dyn ExecutionArgs, + 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, +{ + // The output vector stays at length zero until every slot is initialized so that an unwind + // abandons partially initialized spare capacity. This no-drop assertion proves that no + // initialized value requires a destructor to run. + const { assert_owned_output_needs_no_drop::() }; + + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::const_values(&columns)); + + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let output = &mut values.spare_capacity_mut()[..row_count]; + + let failure = if let Some(views) = Args::views_if_no_consts(&columns) { + // Keep this validation beside the views so LLVM sees their common length here. + vortex_ensure!( + Args::ARITY == 0 || views.len() == row_count, + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: the tuple length check proved every non-nullary view addresses exactly + // `row_count` rows immediately above. Nullary tuples do not access an input view. + let source = unsafe { Args::indexed_source(views, row_count) }; + + source.map_checked_into(output, |elements| apply(&prepared, elements)) + } else { + // Keep this proof branch-local. Shared validation prevents LLVM from specializing this + // loop for each batch-constant arrangement, leaving it scalar under multiple CGUs without + // LTO. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated = Fail::default(); + + // Iterate over `output` directly. A `0..row_count` range reuses the address-taken value + // from the validation error formatter and retains an output bounds check. + for (index, slot) in output.iter_mut().enumerate() { + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the loop. + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + + slot.write(value); + accumulated |= row_failure; + } + + accumulated + }; + + // SAFETY: normal completion of either execution path initializes `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Defer failures so batch execution can retry with only valid rows. + match finish_failure(failure) { + Ok(()) => Ok(RowExecution::Output(Out::build(values))), + Err(error) => Ok(RowExecution::DeferredError(error)), + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs new file mode 100644 index 00000000000..9b809179cae --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -0,0 +1,436 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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. + +use vortex_buffer::BitBuffer; +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::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::ViewLen; + +/// Decode inputs once, then write one sink row for each input row. +/// +/// The executor owns the sink and passes each output row to `apply`. This keeps `apply` as [`Fn`]. +/// Capturing the sink would require [`FnMut`] and put its buffer metadata behind loop-carried +/// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant. +pub(crate) fn execute_sink( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + let columns = Args::decode(args, ctx)?; + + let row_count = args.row_count(); + let const_values = Args::const_values(&columns); + let prepared = prepare(const_values); + + let mut sink = >::with_capacity(row_count)?; + + // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. + { + let mut rows = >::rows(&mut sink); + + // This equality proves to LLVM that `0..row_count` is in bounds for `rows`. + let sink_row_count = rows.len(); + vortex_ensure_eq!( + sink_row_count, + row_count, + "the output sink must address exactly {row_count} rows, got {sink_row_count}", + ); + + let views = Args::views_if_no_consts(&columns); + if let Some(views) = views { + if Args::ARITY != 0 && views.len() != row_count { + decoded_length_error(row_count)?; + } + + for index in 0..row_count { + // SAFETY: the tuple length check proved every non-nullary view has `row_count` + // rows before the loop. Nullary tuples do not access an input view. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + apply(&prepared, elements, output).into_result()?; + } + } else { + if !Args::decoded_lens_match(&columns, row_count) { + decoded_length_error(row_count)?; + } + + for index in 0..row_count { + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the + // loop. + apply(&prepared, Args::get(&columns, index), output).into_result()?; + } + } + } + + // SAFETY: every row callback completed successfully, so each returned the required write token. + unsafe { >::finish(sink) }.map(RowExecution::Output) +} + +/// Write only the rows set in `valid`, or decline when the inputs or sink cannot support +/// skip-invalid execution. +/// +/// `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, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult> +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + let Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + mut sink, + }) = setup_sink_valid_rows::(args, valid, ctx)? + else { + return Ok(None); + }; + + 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`. With multiple CGUs and no LTO, using + // `drop(rows)` duplicates `Args::get` in every sparse callback. + { + // Initialize every slot before visiting only valid rows. + let mut rows = >::rows(&mut 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, + row_count, + "the initialized output sink must address exactly {row_count} rows, got {initialized_row_count}", + ); + + if let Some(views) = views { + if Args::ARITY != 0 && views.len() != row_count { + decoded_length_error(row_count)?; + } + + 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 validated `row_count`. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + // SAFETY: the tuple length check proved every non-nullary view has `row_count` + // rows, and mask indices are below `row_count`. Nullary tuples do not access an + // input view. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + + apply(&prepared, elements, output).into_result() + })?; + } else { + if !Args::decoded_lens_match(&columns, row_count) { + decoded_length_error(row_count)?; + } + + 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 validated `row_count`. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + apply(&prepared, Args::get(&columns, index), output).into_result() + })?; + } + } + + // SAFETY: the initializer completed before traversal, and every visited callback completed + // successfully and returned the required write token. + unsafe { >::finish(sink) } + .map(RowExecution::Output) + .map(Some) +} + +/// Construct a decoded-length error outside the traversal branches. +/// +/// Owned execution (`owned.rs`) derives its index from an output-slice iterator. Sink execution +/// only has indexed row access, so `row_count` remains the loop bound. Formatting the error inside +/// either branch takes the address of that bound and prevents LLVM from vectorizing some sink +/// loops. +#[cold] +#[inline(never)] +fn decoded_length_error(row_count: usize) -> VortexResult<()> { + vortex_bail!("a decoded row input does not address exactly {row_count} rows") +} + +/// State resolved before preparing the skip-invalid row loop. +struct ValidRowsSetup<'valid, Args, Sink, Options> +where + Args: ElementTuple, + Sink: OutputSink, +{ + initialize_skipped_rows: for<'rows> fn(&mut >::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. +fn setup_sink_valid_rows<'valid, Args, Sink, Options>( + args: &dyn ExecutionArgs, + valid: &'valid Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult>> +where + Args: ElementTuple, + Sink: OutputSink, +{ + // The initializer both declares support for skipping rows and initializes those rows. + let Some(initialize_skipped_rows) = >::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 { + return Ok(None); + }; + + let row_count = args.row_count(); + + // Keep allocation before the validity and length checks. With multiple CGUs and no LTO, + // moving it later inlines `Args::get` into every sparse callback, duplicating its bounds + // checks. + let sink = >::with_capacity(row_count)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid_rows) = valid.bit_buffer() else { + vortex_bail!( + "execute_sink_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + vortex_ensure_eq!( + valid_rows.len(), + row_count, + "the validity mask must address exactly {row_count} rows, got {}", + valid_rows.len(), + ); + + Ok(Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + sink, + })) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::RowExecution; + use super::execute_sink_valid_rows; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::NativePType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::InitializedElement; + use crate::scalar_fn::unstable::row::OutputSink; + use crate::scalar_fn::unstable::row::UninitElementSink; + 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 Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn return_dtype(_options: &Options) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize) -> 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. + unsafe impl OutputSink for ShrinkingSink { + type Rows<'a> = &'a mut Vec; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + rows.pop(); + }) + } + + fn return_dtype(_options: &Options) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(vec![0; rows])) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.0 + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::from_iter(self.0).into_array()) + } + } + + #[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 valid = Mask::from_iter([true, false]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( + &args, + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + assert!(execution.is_none()); + + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_initializes_and_writes_addressed_rows() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let valid = Mask::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::< + (i64,), + (), + UninitElementSink, + InitializedElement, + EmptyOptions, + >( + &args, + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + // SAFETY: `output` is the row supplied to this callback. + unsafe { InitializedElement::write(output, value * 2) } + }, + )?; + let Some(RowExecution::Output(actual)) = execution else { + vortex_bail!("the skip-invalid sink must produce an output"); + }; + let expected = PrimitiveArray::from_iter([20_i64, 0, 60]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_rechecks_rows_after_initialization() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( + &args, + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + *output = value; + }, + ); + + let error = match result { + Err(error) => error, + Ok(_) => vortex_bail!("the sink must reject rows changed by its initializer"), + }; + assert!( + error + .to_string() + .contains("initialized output sink must address exactly 2 rows, got 1"), + "unexpected error: {error}", + ); + + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index bcb3a008488..12ce0db87c2 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -11,15 +11,23 @@ //! [`RowFn::dispatch`] implementation uses a [`RowVisitor`] to select an [`ElementTuple`] and //! either an [`OutputElement`] or [`OutputSink`] for each supported dtype combination. //! +//! Unlike a general strict function, a [`RowFn`] cannot produce null from valid inputs. +//! //! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits //! reduce compact failure evidence in that loop and retry only valid rows when null payloads may //! have caused the failure. +// TODO(connor)[RowFn]: Remove this expectation when #9450 connects the batch executor. +#[expect(dead_code)] +mod execute; +pub use execute::RowExecution; + mod row_fn; pub use row_fn::RowFn; mod types; pub use types::ElementTuple; +pub use types::FailureEvidence; pub use types::IndexedElementTuple; pub use types::InitializedElement; pub use types::InputElement; @@ -27,6 +35,7 @@ pub use types::OutputElement; pub use types::OutputSink; pub use types::SinkResult; pub use types::UninitElementSink; +pub use types::ViewLen; mod visitor; pub use visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 8a982c4fb37..831cc25f251 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -19,7 +19,12 @@ use super::visitor::RowVisitor; use crate::dtype::DType; use crate::scalar_fn::ScalarFnId; -/// A scalar function computed one row at a time. +/// A strict scalar function whose row kernel cannot produce null from valid inputs. +/// +/// This is stronger than +/// [`ScalarFnVTable::is_strict`](crate::scalar_fn::ScalarFnVTable::is_strict), which requires null +/// propagation but permits valid inputs to produce null. The framework derives output validity +/// only from input validity. /// /// Declare argument names and use [`dispatch`](Self::dispatch) to select element and output types. /// Every implementation receives the standard [`ScalarFnVTable`]. A public type that needs custom diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs index 28105c893cf..323880148b5 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs @@ -15,7 +15,7 @@ use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::OutputElement; use crate::validity::Validity; -// SAFETY: the per-row view is a bit buffer, and its reported length is the buffer length. +// SAFETY: the view is a bit buffer, and its reported length is the buffer length. unsafe impl InputElement for bool { type Column = BitBuffer; type View<'a> = &'a BitBuffer; @@ -23,7 +23,7 @@ unsafe impl InputElement for bool { // Every bit of the buffer is readable, valid or not. const DENSE_SAFE: bool = true; - const DECODE_FALLIBLE: bool = false; + const DECODE_INFALLIBLE: bool = true; fn validate(dtype: &DType) -> VortexResult<()> { vortex_ensure!( @@ -49,10 +49,6 @@ unsafe impl InputElement for bool { column } - fn view_len(view: &Self::View<'_>) -> usize { - view.len() - } - fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> bool where Self: 'a, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 2764840da7a..bc047721310 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -11,15 +11,18 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; +use crate::scalar_fn::unstable::row::ViewLen; /// An element type that can be read row-wise out of an input column. /// /// # Safety /// -/// For every view returned by [`view`](Self::view), every index below -/// [`view_len`](Self::view_len) **must** satisfy the safety contract of -/// [`get_from_view_unchecked`](Self::get_from_view_unchecked). Shared execution relies on this -/// proof to perform unchecked reads after one pre-loop length check. +/// - For each view returned by [`view`](Self::view), every index below [`ViewLen::len`] **must** +/// satisfy the contract of [`get_from_view_unchecked`](Self::get_from_view_unchecked). +/// - The view length and its addressable indices **must** remain stable while the view exists. +/// Interior mutability exposed through an element **must not** change either property. +/// - Shared execution checks the length once before unchecked reads. Violating these requirements +/// can cause undefined behavior. pub unsafe trait InputElement: 'static { /// The decoded column representation supporting `O(1)` row access. type Column; @@ -29,7 +32,7 @@ pub unsafe trait InputElement: 'static { /// This can borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, /// for example, expose a slice so its pointer and length are loop invariants rather than /// re-reading a [`Buffer`](vortex_buffer::Buffer) descriptor for every row. - type View<'a>; + type View<'a>: ViewLen; /// The borrowed element value handed to a row closure. type Elem<'a>; @@ -41,10 +44,10 @@ pub unsafe trait InputElement: 'static { /// null rows to the row closure. const DENSE_SAFE: bool; - /// Whether [`decode`](Self::decode) can fail on _legal_ input data. + /// Whether [`decode`](Self::decode) is infallible for _legal_ input data. /// /// This excludes infrastructural failures such as IO or allocation. - const DECODE_FALLIBLE: bool; + const DECODE_INFALLIBLE: bool; /// Validate that `dtype` is an acceptable input column dtype for this element type. fn validate(dtype: &DType) -> VortexResult<()>; @@ -68,7 +71,7 @@ pub unsafe trait InputElement: 'static { /// cannot decode this particular array. /// /// Override this for a non-dense-safe representation that can still place safe placeholders in - /// null slots. The skip-invalid executor never reads those slots. + /// null slots. Valid-row execution never reads those slots. fn decode_null_tolerant( array: ArrayRef, ctx: &mut ExecutionCtx, @@ -89,12 +92,6 @@ pub unsafe trait InputElement: 'static { /// keeps their one-row decoded representation separate. fn view(column: &Self::Column) -> Self::View<'_>; - /// Number of rows addressable through a [`View`](Self::View). - /// - /// Every index below this length must be valid for - /// [`get_from_view_unchecked`](Self::get_from_view_unchecked). - fn view_len(view: &Self::View<'_>) -> usize; - /// Read one row from a [`View`](Self::View). fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> where @@ -104,7 +101,7 @@ pub unsafe trait InputElement: 'static { /// /// # Safety /// - /// `index` must be less than [`view_len`](Self::view_len) for `view`. + /// `index` must be less than [`ViewLen::len`] for `view`. unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> where Self: 'a, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs index d1c75b7054a..3ca5af5672c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs @@ -14,7 +14,7 @@ pub trait OutputElement: 'static + Sized { /// derived from the inputs by batch execution. /// /// Because this method takes no arguments, the dtype must be a property of the Rust type. Use - /// an [`OutputSink`] when the output dtype depends on function options or input dtypes. + /// an [`OutputSink`] when the output dtype depends on function options. /// /// [`OutputSink`]: crate::scalar_fn::unstable::row::OutputSink fn element_dtype() -> DType; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs index 9bc5db0c8e4..ed5b7b0888e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs @@ -17,7 +17,7 @@ use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::OutputElement; use crate::validity::Validity; -// SAFETY: the per-row view is a native slice, and its reported length is the slice length. +// SAFETY: the view is a native slice, and its reported length is the slice length. unsafe impl InputElement for T { type Column = Buffer; type View<'a> = &'a [T]; @@ -25,7 +25,7 @@ unsafe impl InputElement for T { // Every lane of the buffer holds a `T`, valid or not. const DENSE_SAFE: bool = true; - const DECODE_FALLIBLE: bool = false; + const DECODE_INFALLIBLE: bool = true; fn validate(dtype: &DType) -> VortexResult<()> { let expected = T::PTYPE; @@ -56,10 +56,6 @@ unsafe impl InputElement for T { column.as_slice() } - fn view_len(view: &Self::View<'_>) -> usize { - view.len() - } - fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> T where Self: 'a, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index b0a3e709696..af42ff56868 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -19,54 +19,63 @@ use crate::arrays::masked::MaskedArraySlotsExt; use crate::dtype::DType; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::ViewLen; /// One decoded input, collapsed to a single row when it is constant for the batch. pub struct ArgColumn( - /// The decoded column, classified by whether it stores one value per row. + /// The decoded argument, classified by how the row loop addresses it. ArgColumnKind, ); enum ArgColumnKind { - PerRow(T::Column), - Constant(T::Column), + /// A decoded column covering the full batch; executors validate its length before traversal. + Column(T::Column), + + /// Exactly one decoded row, established by [`ArgColumn::try_from_const`]. + Const(T::Column), } impl ArgColumn { + fn try_from_const(column: T::Column) -> VortexResult { + let decoded_len = T::view(&column).len(); + vortex_ensure_eq!( + decoded_len, + 1, + "a decoded batch-constant input must contain exactly 1 row, got {decoded_len}", + ); + + Ok(Self(ArgColumnKind::Const(column))) + } + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // An empty input has no row 0 to slice, and its row loop runs zero times either way. - if let Some(constant) = batch_constant(&array) + if let Some(const_array) = batch_const(&array) && !array.is_empty() { - return Ok(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?))); + return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?); } - Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) + Ok(Self(ArgColumnKind::Column(T::decode(array, ctx)?))) } fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { // Batch execution short-circuits null constants before selecting a strategy, so a // constant reaching this path is non-null and can use the ordinary decode. - if let Some(constant) = batch_constant(&array) + if let Some(const_array) = batch_const(&array) && !array.is_empty() { - return Ok(Some(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?)))); + return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?).map(Some); } Ok(T::decode_null_tolerant(array, ctx)? - .map(ArgColumnKind::PerRow) + .map(ArgColumnKind::Column) .map(Self)) } fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { // Batch execution short-circuits null constants before selecting this path, so a // non-empty constant can always use the ordinary decode. - if batch_constant(array).is_some() && !array.is_empty() { + if batch_const(array).is_some() && !array.is_empty() { return Ok(true); } @@ -75,30 +84,30 @@ impl ArgColumn { fn get(&self, index: usize) -> T::Elem<'_> { match &self.0 { - ArgColumnKind::PerRow(column) => T::get(column, index), - ArgColumnKind::Constant(column) => T::get(column, 0), + ArgColumnKind::Column(column) => T::get(column, index), + ArgColumnKind::Const(column) => T::get(column, 0), } } - fn per_row_column(&self) -> Option<&T::Column> { + fn as_column(&self) -> Option<&T::Column> { match &self.0 { - ArgColumnKind::PerRow(column) => Some(column), - ArgColumnKind::Constant(_) => None, + ArgColumnKind::Column(column) => Some(column), + ArgColumnKind::Const(_) => None, } } fn addresses_rows(&self, row_count: usize) -> bool { - // A constant is always read at index zero, so it addresses any batch length. + // A constant is validated when constructed and is always read at index zero. match &self.0 { - ArgColumnKind::PerRow(column) => T::view_len(&T::view(column)) == row_count, - ArgColumnKind::Constant(_) => true, + ArgColumnKind::Column(column) => T::view(column).len() == row_count, + ArgColumnKind::Const(_) => true, } } - fn constant(&self) -> Option> { + fn const_value(&self) -> Option> { match &self.0 { - ArgColumnKind::PerRow(_) => None, - ArgColumnKind::Constant(column) => Some(T::get(column, 0)), + ArgColumnKind::Column(_) => None, + ArgColumnKind::Const(column) => Some(T::get(column, 0)), } } } @@ -107,19 +116,27 @@ impl ArgColumn { /// /// Batch execution owns mask validity, so a masked constant can expose its constant child here. An /// extension over constant storage remains wrapped to preserve its extension dtype. -pub fn batch_constant(array: &ArrayRef) -> Option { +pub fn batch_const(array: &ArrayRef) -> Option { if array.is::() { return Some(array.clone()); } + // TODO(joe): We want to change this to a V2 Masked. if let Some(masked) = array.as_opt::() { - return Some(masked.child().clone()).filter(|child| child.is::()); + return masked + .child() + .is::() + .then(|| masked.child().clone()); + } + + // TODO(connor): This is maybe incorrect unless this is a refinement type? + if let Some(extension) = array.as_opt::() + && extension.storage_array().is::() + { + return Some(array.clone()); } - array - .as_opt::() - .is_some_and(|ext| ext.storage_array().is::()) - .then(|| array.clone()) + None } /// Typed argument tuples for arities zero through twelve. @@ -130,17 +147,17 @@ pub trait ElementTuple: 'static + private::Sealed { /// The decoded column representations. type Columns; - /// Borrowed views of decoded columns when every argument stores one value per row. - type Views<'a>; + /// Borrowed views of decoded columns with no batch constants. + type Views<'a>: ViewLen; /// The borrowed row of element values. type Elems<'a>; /// The batch-constant element values. /// - /// `Some` carries the value of a batch-constant argument. `None` marks a per-row argument. A - /// [`RowVisitor`] passes these values to its prepare closure so constant work can leave the row - /// loop. + /// `Some` carries the value of a batch-constant argument. `None` marks an argument decoded at + /// full batch length. A [`RowVisitor`] passes these values to its prepare closure so constant + /// work can leave the row loop. /// /// [`RowVisitor`]: crate::scalar_fn::unstable::row::RowVisitor type ConstElems<'a>; @@ -151,8 +168,8 @@ pub trait ElementTuple: 'static + private::Sealed { /// Whether every argument is [`InputElement::DENSE_SAFE`]. const DENSE_SAFE: bool; - /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. - const DECODE_FALLIBLE: bool; + /// Whether every argument is [`InputElement::DECODE_INFALLIBLE`]. + const DECODE_INFALLIBLE: bool; /// Validate the input dtypes and exact arity. fn validate(dtypes: &[DType]) -> VortexResult<()>; @@ -170,34 +187,34 @@ pub trait ElementTuple: 'static + private::Sealed { /// Decode every input column once while tolerating null rows. /// - /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid - /// strategy calls this once per batch. + /// Return `Ok(None)` when an argument has no null-tolerant representation. Valid-row execution + /// calls this once per batch. fn decode_null_tolerant( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult>; /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + /// + /// Each argument selects either its batch-constant value or row `index`. Keep that selection + /// visible in the loop so LLVM can unswitch it before vectorizing. fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; - /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// Borrow the decoded columns only when no argument is batch-constant. /// - /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple - /// gives the optimizer ordinary contiguous column access without a per-row constant check. - fn per_row_views(columns: &Self::Columns) -> Option>; + /// Returns `None` if any argument is batch-constant. Otherwise, omitting `ArgColumn` from the + /// returned tuple removes batch-constant checks from the row loop. + fn views_if_no_consts(columns: &Self::Columns) -> Option>; - /// Whether every view contains exactly `row_count` rows. + /// Whether every argument decoded at full batch length contains exactly `row_count` rows. /// - /// The executor calls this once before the all-per-row hot loop. A successful check gives LLVM - /// a dominating equality between the loop bound and every source length, which lets it optimize - /// the tuple access as one fixed-length traversal. - fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; - - /// Whether every per-row argument contains exactly `row_count` rows. + /// NB: `Columns` cannot implement [`ViewLen`] because a batch-constant [`ArgColumn`] physically + /// contains one decoded row while logically addressing every row in the batch. By contrast, + /// [`views_if_no_consts`](Self::views_if_no_consts) constructs `Views` only when every argument + /// is non-constant, so those views have one common physical length. /// - /// This is the mixed-shape equivalent of [`view_lens_match`](Self::view_lens_match) when - /// [`per_row_views`](Self::per_row_views) declines. It runs once before the hot loop for the - /// same LLVM optimization. A batch constant is exempt because decoding collapsed it to one row. + /// This check runs once before the hot loop. A batch constant is exempt because its + /// [`ArgColumn`] constructor already validated the one-row representation produced by decoding. fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; /// Read one row from borrowed views. @@ -213,9 +230,12 @@ pub trait ElementTuple: 'static + private::Sealed { index: usize, ) -> Self::Elems<'a>; - /// Read the batch-constant elements out of the decoded columns once for one row-kernel - /// invocation. - fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; + /// Return one optional batch-constant value per argument. + /// + /// Each tuple position is `Some(value)` when that argument is batch-constant and `None` when it + /// was decoded at full batch length. Executors pass the tuple to the prepare closure once + /// before the row loop. + fn const_values(columns: &Self::Columns) -> Self::ConstElems<'_>; } impl private::Sealed for () {} @@ -228,7 +248,7 @@ impl ElementTuple for () { const ARITY: usize = 0; const DENSE_SAFE: bool = true; - const DECODE_FALLIBLE: bool = false; + const DECODE_INFALLIBLE: bool = true; fn validate(dtypes: &[DType]) -> VortexResult<()> { vortex_ensure_eq!( @@ -257,14 +277,10 @@ impl ElementTuple for () { fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} - fn per_row_views(_columns: &Self::Columns) -> Option> { + fn views_if_no_consts(_columns: &Self::Columns) -> Option> { Some(()) } - fn view_lens_match(_views: &Self::Views<'_>, _row_count: usize) -> bool { - true - } - fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { true } @@ -277,7 +293,7 @@ impl ElementTuple for () { ) -> Self::Elems<'a> { } - fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} + fn const_values(_columns: &Self::Columns) -> Self::ConstElems<'_> {} } macro_rules! element_tuple { @@ -292,7 +308,7 @@ macro_rules! element_tuple { const ARITY: usize = $arity; const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; - const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; + const DECODE_INFALLIBLE: bool = $($t::DECODE_INFALLIBLE &&)+ true; fn validate(dtypes: &[DType]) -> VortexResult<()> { vortex_ensure_eq!( @@ -341,15 +357,8 @@ macro_rules! element_tuple { ($(columns.$idx.get(index),)+) } - fn per_row_views(columns: &Self::Columns) -> Option> { - Some(($($t::view(columns.$idx.per_row_column()?),)+)) - } - - fn view_lens_match( - views: &Self::Views<'_>, - row_count: usize, - ) -> bool { - $($t::view_len(&views.$idx) == row_count &&)+ true + fn views_if_no_consts(columns: &Self::Columns) -> Option> { + Some(($($t::view(columns.$idx.as_column()?),)+)) } fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { @@ -371,8 +380,8 @@ macro_rules! element_tuple { ($(unsafe { $t::get_from_view_unchecked(&views.$idx, index) },)+) } - fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { - ($(columns.$idx.constant(),)+) + fn const_values(columns: &Self::Columns) -> Self::ConstElems<'_> { + ($(columns.$idx.const_value(),)+) } } }; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs index d612d874935..3a21176e25c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -11,13 +11,14 @@ use vortex_compute::lane_kernels::LaneZip; use super::ElementTuple; use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::ViewLen; /// An argument tuple that supports a validated dense indexed traversal. /// /// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's /// unchecked view access after batch execution validates every decoded column length once. pub trait IndexedElementTuple: ElementTuple { - /// The source shared execution uses for a dense all-per-row loop. + /// The source used when no input is batch-constant. /// /// Its length must be the common view length. For every valid index it must preserve row order, /// return the same value as [`ElementTuple::get_from_views`], and uphold the unchecked read @@ -48,7 +49,7 @@ impl<'a, T: InputElement> IndexedSource for ElementSource<'a, T> { type Item = T::Elem<'a>; fn len(&self) -> usize { - T::view_len(&self.view) + self.view.len() } unsafe fn get_unchecked(&self, index: usize) -> Self::Item { diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index 044ad15fd4b..4f8ea67348c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use super::ElementTuple; -use super::element_tuple::batch_constant; +use super::element_tuple::batch_const; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -32,7 +32,7 @@ use crate::validity::Validity; static DECODE_CALLS: AtomicUsize = AtomicUsize::new(0); macro_rules! i64_test_element { - ($element:ident, $decode_fallible:literal $(, $can_decode:item)?) => { + ($element:ident, $decode_infallible:literal $(, $can_decode:item)?) => { struct $element; // SAFETY: the view and unchecked access delegate to the `i64` implementation. @@ -42,7 +42,7 @@ macro_rules! i64_test_element { type Elem<'a> = i64; const DENSE_SAFE: bool = true; - const DECODE_FALLIBLE: bool = $decode_fallible; + const DECODE_INFALLIBLE: bool = $decode_infallible; fn validate(dtype: &DType) -> VortexResult<()> { ::validate(dtype) @@ -63,10 +63,6 @@ macro_rules! i64_test_element { ::view(column) } - fn view_len(view: &Self::View<'_>) -> usize { - ::view_len(view) - } - fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64 where Self: 'a, @@ -87,12 +83,12 @@ macro_rules! i64_test_element { i64_test_element!( DecodeProbe, - false, + true, fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { Ok(true) } ); -i64_test_element!(DenseFallible, true); +i64_test_element!(DenseFallible, false); #[test] fn test_null_tolerant_decline_precedes_decoding() -> VortexResult<()> { @@ -110,31 +106,31 @@ fn test_null_tolerant_decline_precedes_decoding() -> VortexResult<()> { } #[test] -fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { +fn test_batch_const_unwraps_filtered_masked_constant() -> VortexResult<()> { let child = ConstantArray::new(7_i64, 3).into_array(); let masked = MaskedArray::try_new(child, Validity::from_iter([true, false, true]))?.into_array(); let filtered = masked.filter(Mask::from_iter([true, true, false]))?; - let Some(constant) = batch_constant(&filtered) else { + let Some(const_array) = batch_const(&filtered) else { vortex_bail!("filtered masked constant must remain batch-constant"); }; - assert!(constant.is::()); + assert!(const_array.is::()); Ok(()) } #[test] -fn test_batch_constant_preserves_filtered_extension() -> VortexResult<()> { +fn test_batch_const_preserves_filtered_extension() -> VortexResult<()> { let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); let extension = ExtensionArray::new(ext_dtype, ConstantArray::new(7_i64, 3).into_array()).into_array(); let filtered = extension.filter(Mask::from_iter([true, false, true]))?; - let Some(constant) = batch_constant(&filtered) else { + let Some(const_array) = batch_const(&filtered) else { vortex_bail!("filtered extension storage must remain batch-constant"); }; - assert_eq!(constant.dtype(), extension.dtype()); + assert_eq!(const_array.dtype(), extension.dtype()); Ok(()) } 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 ce119f32915..3b4a92e784e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -3,9 +3,10 @@ //! Input decoding and output construction for row functions. //! -//! [`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 -//! the immediate and deferred outcomes returned by sink-writing row closures. +//! [`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. mod element; pub use element::ElementTuple; @@ -14,9 +15,13 @@ pub use element::InputElement; pub use element::OutputElement; mod result; +pub use result::FailureEvidence; pub use result::SinkResult; mod sink; pub use sink::InitializedElement; pub use sink::OutputSink; pub use sink::UninitElementSink; + +mod view; +pub use view::ViewLen; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/result.rs b/vortex-array/src/scalar_fn/unstable/row/types/result.rs index 6163687932d..55c34da0a49 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/result.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/result.rs @@ -1,15 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Return types for sink-writing row closures. +//! Return types for row closures. //! -//! [`SinkResult`] lets the executor handle initialized sinks and sinks that require an -//! [`InitializedElement`] token, with either infallible or immediate-error callbacks. +//! [`FailureEvidence`] represents deferred failures from owned row closures. [`SinkResult`] lets +//! the executor handle initialized sinks and sinks that require an [`InitializedElement`] token, +//! with either infallible or immediate-error callbacks. + +use std::ops::BitOrAssign; use vortex_error::VortexResult; use super::InitializedElement; +/// Compact failure evidence that can be OR-reduced across rows. +/// +/// [`Default::default`] **must** mean success, including for an empty batch. The compiler cannot +/// check this requirement. +pub trait FailureEvidence: Copy + Default + BitOrAssign {} + +impl FailureEvidence for T {} + /// The result of writing one row: success or an immediate error. /// /// This trait is sealed. Row functions choose one of its supplied implementations. 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 6cc3ce06d30..e95431bf9e6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -14,15 +14,15 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::dtype::DType; use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::ViewLen; /// A column allocated once per batch that a row closure writes into, one row at a time. /// -/// A sink can use function options and input dtypes to build a runtime-shaped output or own shared -/// batch state. The executor passes each row slot into an [`Fn`] closure. +/// A sink can use function options to build a runtime-shaped output or own shared batch state. The +/// executor passes each row slot into an [`Fn`] closure. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. -/// Skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an -/// initializer. +/// Execution can omit invalid rows when [`skipped_rows_initializer`] returns an initializer. /// /// # Errors /// @@ -34,21 +34,26 @@ use crate::scalar_fn::unstable::row::OutputElement; /// /// An implementation must uphold all of these requirements: /// -/// - Every index in `0..row_count(rows)` **must** identify one distinct row owned by this sink. +/// - Every index below [`ViewLen::len`] for [`Rows`] **must** identify one distinct row owned by +/// this sink. +/// - 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 +/// 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 /// an uninitialized row **must not** be safely forgeable, reusable, or substitutable. -/// - An initializer returned by [`skipped_rows_initializer`] **must** initialize every row. /// - `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. +/// - Violating these requirements can cause undefined behavior. /// /// [`Rows`]: Self::Rows /// [`WriteToken`]: Self::WriteToken /// [`finish`]: Self::finish -/// [`row_count`]: Self::row_count /// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE /// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult /// [`skipped_rows_initializer`]: Self::skipped_rows_initializer @@ -57,7 +62,7 @@ pub unsafe trait OutputSink: 'static + Sized { /// /// 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> + type Rows<'a>: ViewLen where Self: 'a; @@ -73,21 +78,22 @@ 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 skip-invalid execution. + /// The operation that initializes every output position before + /// [skip-invalid execution](crate::scalar_fn::unstable::row). /// - /// `Some` enables skip-invalid execution. The initializer **must** make every row safe to - /// finish. Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// `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 the executor fall back to filtering the inputs. fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { None } - /// The dtype of the column this sink builds, given the function options and input dtypes. + /// The dtype of the column this sink builds, given the function options. /// /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the /// result, and masks the null rows. - fn output_dtype(options: &Options, args: &[DType]) -> VortexResult; + fn return_dtype(options: &Options) -> VortexResult; /// Allocate a sink for `rows` rows. fn with_capacity(rows: usize) -> VortexResult; @@ -95,18 +101,15 @@ pub unsafe trait OutputSink: 'static + Sized { /// Borrow all output rows for the hot loop. fn rows(&mut self) -> Self::Rows<'_>; - /// The number of rows addressable through [`row_unchecked`](Self::row_unchecked). - fn row_count(rows: &Self::Rows<'_>) -> usize; - /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. /// /// # Safety /// - /// `index` must be less than [`row_count`](Self::row_count) for `rows`. + /// `index` must be less than [`ViewLen::len`] for `rows`. unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; /// Finish into the built column, whose dtype **must** be this sink's - /// [`output_dtype`](Self::output_dtype). Called once per batch. + /// [`return_dtype`](Self::return_dtype). Called once per batch. /// /// # Safety /// @@ -153,7 +156,7 @@ impl InitializedElement { /// The row closure must return the [`InitializedElement`] from [`InitializedElement::write`] on /// success. The token is zero-sized, so the proof adds no runtime row state. /// -/// Skip-invalid execution initializes placeholders before omitting rows. Errors and unwinds are +/// When execution omits invalid rows, it initializes placeholders first. Errors and unwinds are /// safe because `values` keeps length zero until `finish`. The `T: Copy` bound means that /// initialized spare-capacity elements require no destruction. pub struct UninitElementSink { @@ -183,7 +186,7 @@ unsafe impl OutputSink }) } - fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn return_dtype(_options: &Options) -> VortexResult { Ok(T::element_dtype()) } @@ -198,10 +201,6 @@ unsafe impl OutputSink &mut self.values.spare_capacity_mut()[..self.row_count] } - fn row_count(rows: &Self::Rows<'_>) -> usize { - rows.len() - } - 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) } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/view.rs b/vortex-array/src/scalar_fn/unstable/row/types/view.rs new file mode 100644 index 00000000000..18085a12d14 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/view.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Length reporting for row-loop views. +//! +//! [`ViewLen`] lets input and output abstractions expose the number of addressable rows through +//! their borrowed view types. + +use vortex_buffer::BitBuffer; + +/// The number of rows addressable through a row-loop view. +pub trait ViewLen { + /// Return the number of addressable rows. + fn len(&self) -> usize; + + /// Return whether the view contains no rows. + fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +impl ViewLen for () { + fn len(&self) -> usize { + 0 + } +} + +impl ViewLen for BitBuffer { + fn len(&self) -> usize { + BitBuffer::len(self) + } +} + +impl ViewLen for [T] { + fn len(&self) -> usize { + <[T]>::len(self) + } +} + +impl ViewLen for Vec { + fn len(&self) -> usize { + Vec::len(self) + } +} + +impl ViewLen for &T { + fn len(&self) -> usize { + T::len(self) + } +} + +impl ViewLen for &mut T { + fn len(&self) -> usize { + T::len(self) + } +} + +macro_rules! impl_tuple_view_len { + ($first:ident; $($rest:ident : $idx:tt),*) => { + impl<$first: ViewLen, $($rest: ViewLen),*> ViewLen for ($first, $($rest,)*) { + fn len(&self) -> usize { + let len = self.0.len(); + $(assert_eq!(self.$idx.len(), len, "tuple views must have equal lengths");)* + + len + } + } + }; +} + +impl_tuple_view_len!(A;); +impl_tuple_view_len!(A; B: 1); +impl_tuple_view_len!(A; B: 1, C: 2); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4, F: 5); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4, F: 5, G: 6); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10); +impl_tuple_view_len!(A; B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10, L: 11); + +#[cfg(test)] +mod tests { + use super::ViewLen; + + #[test] + fn tuple_len_returns_common_len() { + let first: &[i64] = &[1, 2]; + let second: &[i64] = &[3, 4]; + + assert_eq!((first, second).len(), 2); + } + + #[test] + #[should_panic(expected = "tuple views must have equal lengths")] + fn tuple_len_rejects_mismatch() { + let first: &[i64] = &[1]; + let second: &[i64] = &[2, 3]; + + let _ = (first, second).len(); + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs index 12d5bd7c7ea..7606f82f0b5 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -7,13 +7,13 @@ //! selected visit with the input dtypes during planning and return its output dtype. use std::mem::needs_drop; -use std::ops::BitOrAssign; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::dtype::DType; use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; @@ -36,7 +36,7 @@ const fn assert_input_visit_contract() { // Dictionary push-down can evaluate values that no input row references. Every dispatch must // therefore match the function-wide fallibility declaration. assert!( - !Args::DECODE_FALLIBLE || F::FALLIBLE, + Args::DECODE_INFALLIBLE || F::FALLIBLE, "RowFn::FALLIBLE must be true when input decoding can fail", ); } @@ -69,7 +69,7 @@ where Function: RowFn, Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + Fail: FailureEvidence, { assert_owned_visit_contract::(); assert!( @@ -106,7 +106,7 @@ where { Args::validate(dtypes)?; - let dtype = Sink::output_dtype(options, dtypes)?; + let dtype = Sink::return_dtype(options)?; vortex_ensure!( !dtype.is_nullable(), "row output sinks must declare a non-nullable dtype, got {dtype}", diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs new file mode 100644 index 00000000000..610568343f3 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -0,0 +1,305 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visitors that execute dense and skip-invalid row loops. +//! +//! Each visit revalidates its concrete signature and checks that its output dtype and null 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; +use vortex_mask::Mask; + +use super::RowPolicy; +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::row_visitor::private; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::FailureEvidence; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +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::RowExecution; +use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_sink; +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. +pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { + /// The inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, + output_dtype: &'args DType, + policy: RowPolicy, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + options, + output_dtype, + policy, + ctx, + } + } +} + +impl private::Sealed for ExecuteRows<'_, '_, F> {} + +impl RowVisitor for ExecuteRows<'_, '_, F> { + type VisitResult = RowExecution; + + 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::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + execute_owned_infallible::(self.args, self.ctx, prepare, apply) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink::( + self.args, 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::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; + + execute_owned::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} + +/// The runtime visit that executes valid rows over the original input columns. +/// +/// Only output sinks have a contract for skipped output positions. Owned visits therefore decline +/// so batch execution can filter the valid inputs and scatter the output back. +pub(crate) struct ExecuteValidRows<'args, 'ctx, F: RowFn> { + /// The original inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The conjoined validity, containing both valid and invalid rows. + valid: &'args Mask, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, + output_dtype: &'args DType, + policy: RowPolicy, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + options, + output_dtype, + policy, + valid, + ctx, + } + } +} + +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} + +impl RowVisitor for ExecuteValidRows<'_, '_, F> { + type VisitResult = Option; + + 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::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + // Owned execution has no sink that can initialize skipped output positions. Decline so + // batch execution filters the inputs and retries with the dense visitor. + Ok(None) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink_valid_rows::( + self.args, self.valid, 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::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; + + // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. + Ok(None) + } +} + +fn ensure_plan( + planned_output: &DType, + planned_policy: RowPolicy, + actual_output: DType, + actual_policy: RowPolicy, +) -> VortexResult<()> { + vortex_ensure_eq!( + actual_policy, + planned_policy, + "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", + ); + vortex_ensure_eq!( + actual_output, + *planned_output, + "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", + ); + + Ok(()) +} 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 57da5f4691b..1b063c788db 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -6,9 +6,15 @@ //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch mod check; +pub(super) use check::assert_owned_output_needs_no_drop; + +// TODO(connor)[RowFn]: Remove this expectation when #9450 constructs the execution visitors. +#[expect(dead_code)] +mod execute; mod plan; pub(super) use plan::BatchPlanner; +pub(super) use plan::RowPolicy; mod row_visitor; pub use row_visitor::RowVisitor; 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 c522376d491..fe3c5e32bb3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -7,7 +7,6 @@ //! null-handling policy that execution must reproduce. use std::marker::PhantomData; -use std::ops::BitOrAssign; use vortex_error::VortexResult; @@ -21,6 +20,7 @@ use super::row_visitor::private; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; @@ -98,7 +98,7 @@ impl RowVisitor for BatchPlanner<'_, F> { where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + Fail: FailureEvidence, { const { assert_deferred_visit_contract::() }; Ok(BatchPlan { @@ -114,15 +114,14 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. - // TODO(connor)[RowFn]: Remove this allowance when the execution backend from #9130 consumes - // this policy. - #[allow(dead_code)] + // TODO(connor)[RowFn]: Remove this expectation when #9450 uses the planned policy. + #[expect(dead_code)] pub(crate) policy: RowPolicy, } impl BatchPlan { /// Return the output dtype widened with strict input nullability. - pub(crate) fn result_dtype(self, args: &[DType]) -> DType { + pub(crate) fn result_dtype(&self, args: &[DType]) -> DType { let nullability = self.output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); @@ -139,14 +138,14 @@ pub(crate) enum RowPolicy { /// Evaluate all rows, retrying only valid rows if a deferred error is raised. DenseWithRetry, - /// Execute only valid rows, trying skip-invalid execution before filtering. + /// Execute only valid rows over the original inputs before filtering. ValidOnly, } impl RowPolicy { /// The policy for an infallible owned output. pub(crate) const fn for_owned_output() -> Self { - if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + if Args::DENSE_SAFE && Args::DECODE_INFALLIBLE { Self::Dense } else { Self::ValidOnly @@ -155,7 +154,7 @@ impl RowPolicy { /// The policy for an owned output carrying batch-deferred failure evidence. pub(crate) const fn for_deferred_output() -> Self { - if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + if Args::DENSE_SAFE && Args::DECODE_INFALLIBLE { Self::DenseWithRetry } else { Self::ValidOnly @@ -164,7 +163,7 @@ impl RowPolicy { /// The policy for a sink-writing output. pub(crate) const fn for_sink() -> Self { - if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { + if Args::DENSE_SAFE && Args::DECODE_INFALLIBLE && !ApplyResult::FALLIBLE { Self::Dense } else { Self::ValidOnly diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 0e7abaa4322..ce46e53210b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -8,11 +8,10 @@ //! //! [`RowFn`]: crate::scalar_fn::unstable::row::RowFn -use std::ops::BitOrAssign; - use vortex_error::VortexResult; use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; @@ -74,7 +73,8 @@ pub trait RowVisitor: private::Sealed + Sized { /// # Examples /// /// Test whether each string occurs in its allowed-values list. The prepare closure builds one - /// lookup table for a batch-constant list. The row closure scans a varying list directly. + /// lookup table for a batch-constant list. The row closure scans the current list from the input + /// column directly. /// /// ```ignore /// visitor.visit_prepared::< @@ -161,16 +161,16 @@ pub trait RowVisitor: private::Sealed + Sized { /// visitor.visit_prepared_into::< /// (TensorRow, TensorRow), /// UninitElementSink, - /// ConstantVectorMagnitudes, + /// ConstVectorMagnitudes, /// InitializedElement, /// >( - /// |(lhs, rhs)| ConstantVectorMagnitudes { + /// |(lhs, rhs)| ConstVectorMagnitudes { /// lhs: lhs.map(vector_magnitude), /// rhs: rhs.map(vector_magnitude), /// }, - /// |constant_magnitudes, (lhs, rhs), output| { + /// |const_magnitudes, (lhs, rhs), output| { /// let similarity = - /// cosine_similarity_with_constant_magnitudes(constant_magnitudes, lhs, rhs); + /// cosine_similarity_with_const_magnitudes(const_magnitudes, lhs, rhs); /// /// // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. /// unsafe { InitializedElement::write(output, similarity) } @@ -196,9 +196,8 @@ pub trait RowVisitor: private::Sealed + Sized { /// `apply` must not panic or have side effects. Dense execution can pass unspecified values /// from null rows. /// - /// The executor OR-reduces `Fail` across rows and passes the result to `finish_failure`. - /// [`Default::default`] **must** mean success, including for an empty batch. The compiler - /// cannot check this requirement. + /// The executor OR-reduces [`FailureEvidence`] across rows and passes the result to + /// `finish_failure`. /// /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) **must** be `true`. /// `Out` must not require drop glue. `Fail` must be no wider than `Out`, or failure tracking @@ -237,7 +236,7 @@ pub trait RowVisitor: private::Sealed + Sized { where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + Fail: FailureEvidence, { self.visit_prepared_deferred::( |_| (), @@ -290,7 +289,7 @@ pub trait RowVisitor: private::Sealed + Sized { where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign; + Fail: FailureEvidence; } pub(super) mod private { diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index aac8ead42fe..1da0b037f2a 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -479,6 +479,33 @@ impl BitBuffer { } } + /// Fallible variant of [`for_each_set_index`](Self::for_each_set_index). + /// + /// Stops and returns the first error from `f`. + #[inline] + pub fn try_for_each_set_index(&self, mut f: F) -> Result<(), E> + where + F: FnMut(usize) -> Result<(), E>, + { + let mut base = 0usize; + for word in self.chunks().iter_padded() { + if word == u64::MAX { + for k in 0..64 { + f(base + k)?; + } + } else { + let mut w = word; + while w != 0 { + f(base + w.trailing_zeros() as usize)?; + w &= w - 1; + } + } + base += 64; + } + + Ok(()) + } + /// Created a new BitBuffer with offset reset to 0 pub fn sliced(&self) -> Self { if self.offset.is_multiple_of(8) { @@ -970,12 +997,21 @@ mod tests { #[case(65)] #[case(200)] #[case(1000)] - fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + fn test_set_index_visitors_match_set_indices(#[case] len: usize) { let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); buf.for_each_set_index(|i| got.push(i)); assert_eq!(got, expected); + + let mut fallible_got = Vec::new(); + let result = buf.try_for_each_set_index(|i| { + fallible_got.push(i); + Ok::<(), ()>(()) + }); + assert_eq!(result, Ok(())); + assert_eq!(fallible_got, expected); } #[rstest] @@ -998,6 +1034,33 @@ mod tests { assert_eq!(got, (0..130).collect::>()); } + #[test] + fn test_try_for_each_set_index_stops_on_error() { + for (buffer, stop) in [ + (BitBuffer::new_set(130), 65), + (BitBuffer::collect_bool(130, |i| i % 3 == 0), 66), + ] { + let mut visited = Vec::new(); + let result = buffer.try_for_each_set_index(|index| { + visited.push(index); + if index == stop { + return Err(index); + } + + Ok(()) + }); + + assert_eq!(result, Err(stop)); + assert_eq!( + visited, + buffer + .set_indices() + .take_while(|&i| i <= stop) + .collect::>() + ); + } + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value