From e7e7c830474cf78b14ef0a6b0192c1429b0ce87c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 17 Aug 2026 15:05:06 -0400 Subject: [PATCH 1/5] Implement RowFn row execution Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/execute/mod.rs | 19 + .../scalar_fn/unstable/row/execute/outcome.rs | 43 ++ .../scalar_fn/unstable/row/execute/owned.rs | 128 ++++++ .../scalar_fn/unstable/row/execute/sink.rs | 380 ++++++++++++++++++ .../src/scalar_fn/unstable/row/mod.rs | 5 + .../src/scalar_fn/unstable/row/row_fn.rs | 7 +- .../unstable/row/types/element/input.rs | 2 +- .../unstable/row/types/element/mod.rs | 1 + .../row/types/element/tuple/element_tuple.rs | 41 +- .../unstable/row/types/element/tuple/mod.rs | 1 + .../unstable/row/types/element/tuple/tests.rs | 2 +- .../src/scalar_fn/unstable/row/types/mod.rs | 1 + .../src/scalar_fn/unstable/row/types/sink.rs | 12 +- .../scalar_fn/unstable/row/visitor/execute.rs | 306 ++++++++++++++ .../src/scalar_fn/unstable/row/visitor/mod.rs | 7 + .../scalar_fn/unstable/row/visitor/plan.rs | 7 +- .../unstable/row/visitor/row_visitor.rs | 2 +- 17 files changed, 933 insertions(+), 31 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/owned.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/sink.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs 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..e3e64194cae --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -0,0 +1,128 @@ +// 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::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; + +/// Zero-sized evidence used to erase failure reduction from infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column for one kernel invocation, then store one infallible output per row. +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 for one kernel invocation, then store outputs and reduce failures. +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: Copy + Default + BitOrAssign, +{ + const { assert_owned_output_needs_no_drop::() }; + + // Keep the vector length at zero until every row succeeds. An unwind then abandons partially + // initialized spare capacity without treating it as initialized output. The no-drop assertion + // above proves that no initialized value requires its destructor to run. + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let failure; + + { + let output = &mut values.spare_capacity_mut()[..row_count]; + + // When every input stores one value per row, the indexed source removes argument-shape + // dispatch from the hot loop and lets the lane kernel optimize the traversal as one + // operation. Keep view construction and its length proof in this branch. Hoisting them + // through the shared validation helper changed add, subtract, and multiply with + // batch-constant and per-row arguments from 9.219, 9.229, and 18.94 us to 30.46, 31.11, + // and 37.73 us on a Ryzen 9 7950X with rustc 1.91.0 and LLVM 21.1.2. + // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. + if let Some(views) = Args::per_row_views(&columns) { + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + failure = unsafe { Args::indexed_source(views, row_count) } + .map_checked_into(output, |elements| apply(&prepared, elements)); + } else { + // A batch-constant input was collapsed to one row during decoding. This path reads that + // row repeatedly while indexing only the per-row inputs. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // Keep the output-slot iterator as the loop bound. `row_count` is address-taken by the + // validation error formatting above. With rustc 1.97.1 and LLVM 22.1.6 under 16 CGUs + // without LTO, indexing `output` by a `0..row_count` range retains an early-exit bounds + // check and prevents vectorization with batch-constant and per-row arguments. Recheck + // the optimized IR and those benchmarks before restoring that range loop. + let mut accumulated = Fail::default(); + for (index, slot) in output.iter_mut().enumerate() { + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + slot.write(value); + accumulated |= row_failure; + } + failure = accumulated; + } + } + + // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Failure evidence is reduced inside the loop so its richer error construction stays cold. + // Preserve that provenance so batch execution may retry over 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..ec121b21b44 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -0,0 +1,380 @@ +// 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 can instead initialize omitted output +//! positions and visit only rows that are valid in every input, falling back when either the input +//! representation or sink lacks that capability. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +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; + +fn ensure_decoded_lengths( + columns: &Args::Columns, + views: Option<&Args::Views<'_>>, + row_count: usize, +) -> VortexResult<()> { + let lengths_match = match views { + Some(views) => Args::view_lens_match(views, row_count), + None => Args::decoded_lens_match(columns, row_count), + }; + vortex_ensure!( + lengths_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(()) +} + +/// Decode every input column and allocate one sink for one kernel invocation. +/// +/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state +/// does not need to be captured by the closure. +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 row_count = args.row_count(); + let mut sink = >::with_capacity(row_count)?; + let columns = Args::decode(args, ctx)?; + let constants = Args::constants(&columns); + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + let prepared = prepare(constants); + + { + // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This + // scope releases the borrow before `finish_sink` consumes the sink. + let mut rows = >::rows(&mut sink); + let sink_row_count = >::row_count(&rows); + vortex_ensure_eq!( + sink_row_count, + row_count, + "the output sink must address exactly {row_count} rows, got {sink_row_count}", + ); + + // The all-per-row representation removes argument-shape dispatch from the hot loop. The + // constant-and-per-row path instead reads collapsed batch constants at row zero. + if let Some(views) = views { + for index in 0..row_count { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before + // the loop. + 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 { + 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) }; + apply(&prepared, Args::get(&columns, index), output).into_result()?; + } + } + } + + finish_sink::(sink) +} + +/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +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>, +{ + // Decline before input decoding or sink allocation when this sink cannot initialize rows that + // the mask skips. The capability and the operation are the same function pointer. + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An + // element representation may decline when it cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + let constants = Args::constants(&columns); + let row_count = args.row_count(); + let mut sink = >::with_capacity(row_count)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid) = 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.len(), + row_count, + "the validity mask must address exactly {row_count} rows, got {}", + valid.len(), + ); + + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + let prepared = prepare(constants); + + { + let mut rows = >::rows(&mut sink); + + // Initialize every slot before skipping rows. Recheck addressability afterward because the + // initializer mutably borrows the row representation. + initialize_skipped_rows(&mut rows); + let initialized_row_count = >::row_count(&rows); + vortex_ensure_eq!( + initialized_row_count, + row_count, + "the initialized output sink must address exactly {row_count} rows, got {initialized_row_count}", + ); + + // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first + // immediate error, turn later callbacks into no-ops, and return before finishing the sink. + let mut error = None; + valid.for_each_set_index(|index| { + if error.is_some() { + return; + } + + // 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) }; + let result = match &views { + Some(views) => { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows, and + // mask indices are below `row_count`. + let elements = unsafe { Args::get_from_views_unchecked(views, index) }; + apply(&prepared, elements, output) + } + None => apply(&prepared, Args::get(&columns, index), output), + }; + if let Err(row_error) = result.into_result() { + error = Some(row_error); + } + }); + + if let Some(error) = error { + return Err(error); + } + } + + finish_sink::(sink).map(Some) +} + +fn finish_sink(sink: Sink) -> VortexResult +where + Sink: OutputSink, +{ + // SAFETY: callers reach this helper only after every completed callback returned the sink's + // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. + // The sink contract defines how that evidence establishes initialization of its row storage. + unsafe { >::finish(sink) }.map(RowExecution::Output) +} + +#[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 output_dtype(_options: &Options, _args: &[DType]) -> 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<'_> {} + + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 + } + + 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 output_dtype(_options: &Options, _args: &[DType]) -> 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 + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + 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..057b5cda694 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -11,10 +11,15 @@ //! [`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. +mod execute; +pub use execute::RowExecution; + mod row_fn; pub use row_fn::RowFn; 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/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 2764840da7a..c9140c00b8c 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 @@ -68,7 +68,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, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 51d66594332..7a8a8e92e41 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -20,3 +20,4 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; +pub use tuple::batch_constant; 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..e6f63628e88 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 @@ -27,20 +27,31 @@ pub struct ArgColumn( ); enum ArgColumnKind { + /// One decoded value per batch row; executors validate the exact length before traversal. PerRow(T::Column), + + /// Exactly one decoded row, established by [`ArgColumn::try_from_constant`]. Constant(T::Column), } impl ArgColumn { + fn try_from_constant(column: T::Column) -> VortexResult { + let decoded_len = T::view_len(&T::view(&column)); + vortex_ensure_eq!( + decoded_len, + 1, + "a decoded batch-constant input must contain exactly 1 row, got {decoded_len}", + ); + + Ok(Self(ArgColumnKind::Constant(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) && !array.is_empty() { - return Ok(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?); } Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) @@ -52,10 +63,7 @@ impl ArgColumn { if let Some(constant) = batch_constant(&array) && !array.is_empty() { - return Ok(Some(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?)))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?).map(Some); } Ok(T::decode_null_tolerant(array, ctx)? @@ -88,14 +96,14 @@ impl ArgColumn { } 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, } } - fn constant(&self) -> Option> { + fn constant_value(&self) -> Option> { match &self.0 { ArgColumnKind::PerRow(_) => None, ArgColumnKind::Constant(column) => Some(T::get(column, 0)), @@ -170,8 +178,8 @@ 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, @@ -195,9 +203,10 @@ pub trait ElementTuple: 'static + private::Sealed { /// Whether every per-row argument contains exactly `row_count` rows. /// - /// 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 is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include + /// batch constants. It runs once before the hot loop for the same LLVM optimization. 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. @@ -372,7 +381,7 @@ macro_rules! element_tuple { } fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { - ($(columns.$idx.constant(),)+) + ($(columns.$idx.constant_value(),)+) } } }; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index a2c143704a0..69b5cf686f6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -8,6 +8,7 @@ mod element_tuple; pub use element_tuple::ElementTuple; +pub use element_tuple::batch_constant; mod indexed; pub use indexed::IndexedElementTuple; 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..560e1e48f4f 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::batch_constant; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; 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..e47f195410c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -12,6 +12,7 @@ pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; +pub(super) use element::batch_constant; mod result; pub use result::SinkResult; 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..22c90b8da02 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -21,8 +21,7 @@ use crate::scalar_fn::unstable::row::OutputElement; /// 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 /// @@ -73,10 +72,11 @@ 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>)> { @@ -153,7 +153,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 { 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..b003f70a395 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -0,0 +1,306 @@ +// 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 std::ops::BitOrAssign; + +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::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: Copy + Default + BitOrAssign, + { + 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: Copy + Default + BitOrAssign, + { + 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..c7f9baf6a62 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,16 @@ //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch mod check; +pub(super) use check::assert_owned_output_needs_no_drop; + +mod execute; +pub(super) use execute::ExecuteRows; +pub(super) use execute::ExecuteValidRows; mod plan; +pub(super) use plan::BatchPlan; 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..5360ced573b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -114,15 +114,12 @@ 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)] 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,7 +136,7 @@ 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, } 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..1a9e8628bc7 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 @@ -74,7 +74,7 @@ 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 a per-row list directly. /// /// ```ignore /// visitor.visit_prepared::< From eb4a7c050529bfc925a288062c9ecf2f29e7ae05 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 17 Aug 2026 16:49:33 -0400 Subject: [PATCH 2/5] Refine RowFn row execution Signed-off-by: Connor Tsui --- .../scalar_fn/unstable/row/execute/owned.rs | 101 ++++---- .../scalar_fn/unstable/row/execute/sink.rs | 219 +++++++++++------- .../row/types/element/tuple/element_tuple.rs | 31 +-- .../row/types/element/tuple/indexed.rs | 2 +- vortex-buffer/src/bit/buf.rs | 65 +++++- 5 files changed, 269 insertions(+), 149 deletions(-) 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 e3e64194cae..d3809f939be 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -20,7 +20,7 @@ use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; -/// Zero-sized evidence used to erase failure reduction from infallible owned visits. +/// Zero-sized failure accumulator for infallible owned visits. #[derive(Clone, Copy, Default)] struct NoFailure; @@ -28,7 +28,7 @@ impl BitOrAssign for NoFailure { fn bitor_assign(&mut self, _rhs: Self) {} } -/// Decode every input column for one kernel invocation, then store one infallible output per row. +/// 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, @@ -48,7 +48,7 @@ where ) } -/// Decode every input column for one kernel invocation, then store outputs and reduce failures. +/// Decode every input column, then store outputs and combine per-row failure evidence. pub(crate) fn execute_owned( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, @@ -61,66 +61,59 @@ where Out: OutputElement, Fail: Copy + Default + BitOrAssign, { + // 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::() }; - // Keep the vector length at zero until every row succeeds. An unwind then abandons partially - // initialized spare capacity without treating it as initialized output. The no-drop assertion - // above proves that no initialized value requires its destructor to run. - let row_count = args.row_count(); - let mut values = Vec::::with_capacity(row_count); let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); - let failure; - - { - let output = &mut values.spare_capacity_mut()[..row_count]; - - // When every input stores one value per row, the indexed source removes argument-shape - // dispatch from the hot loop and lets the lane kernel optimize the traversal as one - // operation. Keep view construction and its length proof in this branch. Hoisting them - // through the shared validation helper changed add, subtract, and multiply with - // batch-constant and per-row arguments from 9.219, 9.229, and 18.94 us to 30.46, 31.11, - // and 37.73 us on a Ryzen 9 7950X with rustc 1.91.0 and LLVM 21.1.2. - // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. - if let Some(views) = Args::per_row_views(&columns) { - vortex_ensure!( - Args::view_lens_match(&views, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - - // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows - // immediately above. - failure = unsafe { Args::indexed_source(views, row_count) } - .map_checked_into(output, |elements| apply(&prepared, elements)); - } else { - // A batch-constant input was collapsed to one row during decoding. This path reads that - // row repeatedly while indexing only the per-row inputs. - vortex_ensure!( - Args::decoded_lens_match(&columns, row_count), - "a decoded row input does not address exactly {row_count} rows", - ); - - // Keep the output-slot iterator as the loop bound. `row_count` is address-taken by the - // validation error formatting above. With rustc 1.97.1 and LLVM 22.1.6 under 16 CGUs - // without LTO, indexing `output` by a `0..row_count` range retains an early-exit bounds - // check and prevents vectorization with batch-constant and per-row arguments. Recheck - // the optimized IR and those benchmarks before restoring that range loop. - let mut accumulated = Fail::default(); - for (index, slot) in output.iter_mut().enumerate() { - let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); - slot.write(value); - accumulated |= row_failure; - } - failure = accumulated; + + 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_no_constants(&columns) { + // Keep this validation beside the views so LLVM sees their common length here. + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + 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. The exact pass interaction is unknown. + 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; } - } - // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + 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) }; - // Failure evidence is reduced inside the loop so its richer error construction stays cold. - // Preserve that provenance so batch execution may retry over only valid rows. + // 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 index ec121b21b44..4fd2b42476b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -3,10 +3,11 @@ //! Executes row kernels that write through an [`OutputSink`]. //! -//! Dense execution visits every row. Skip-invalid execution can instead initialize omitted output -//! positions and visit only rows that are valid in every input, falling back when either the input -//! representation or sink lacks that capability. +//! 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; @@ -21,7 +22,12 @@ use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::SinkResult; -fn ensure_decoded_lengths( +/// Verify that every decoded input addresses exactly `row_count` rows. +/// +/// Unlike the owned executor, the paths with and without batch constants can share this check +/// without losing sink-loop vectorization under multiple CGUs without LTO. The exact pass +/// interaction is unknown. +fn verify_lengths( columns: &Args::Columns, views: Option<&Args::Views<'_>>, row_count: usize, @@ -30,6 +36,7 @@ fn ensure_decoded_lengths( Some(views) => Args::view_lens_match(views, row_count), None => Args::decoded_lens_match(columns, row_count), }; + vortex_ensure!( lengths_match, "a decoded row input does not address exactly {row_count} rows", @@ -38,10 +45,11 @@ fn ensure_decoded_lengths( Ok(()) } -/// Decode every input column and allocate one sink for one kernel invocation. +/// Decode inputs once, then write one sink row for each input row. /// -/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state -/// does not need to be captured by the closure. +/// 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, @@ -53,18 +61,22 @@ where Sink: OutputSink, ApplyResult: SinkResult>::WriteToken>, { - let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count)?; let columns = Args::decode(args, ctx)?; + let views = Args::views_no_constants(&columns); + + let row_count = args.row_count(); + verify_lengths::(&columns, views.as_ref(), row_count)?; + let constants = Args::constants(&columns); - let views = Args::per_row_views(&columns); - ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; let prepared = prepare(constants); + let mut sink = >::with_capacity(row_count)?; + + // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. { - // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This - // scope releases the borrow before `finish_sink` 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 = >::row_count(&rows); vortex_ensure_eq!( sink_row_count, @@ -72,16 +84,14 @@ where "the output sink must address exactly {row_count} rows, got {sink_row_count}", ); - // The all-per-row representation removes argument-shape dispatch from the hot loop. The - // constant-and-per-row path instead reads collapsed batch constants at row zero. if let Some(views) = views { for index in 0..row_count { - // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before - // the loop. + // SAFETY: `verify_lengths` proved every view has `row_count` rows before the loop. 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 { @@ -89,15 +99,23 @@ where // 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()?; } } } - finish_sink::(sink) + // SAFETY: every row callback completed successfully, so each returned the required write token. + unsafe { >::finish(sink) }.map(RowExecution::Output) } -/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +/// 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, @@ -110,45 +128,32 @@ where Sink: OutputSink, ApplyResult: SinkResult>::WriteToken>, { - // Decline before input decoding or sink allocation when this sink cannot initialize rows that - // the mask skips. The capability and the operation are the same function pointer. - let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + let Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + mut sink, + }) = setup_sink_valid_rows::(args, valid, ctx)? else { return Ok(None); }; - // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An - // element representation may decline when it cannot provide those values safely. - let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { - return Ok(None); - }; - let constants = Args::constants(&columns); - let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count)?; - - // Batch execution resolves all-valid and all-null inputs before selecting this path. - let AllOr::Some(valid) = 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.len(), - row_count, - "the validity mask must address exactly {row_count} rows, got {}", - valid.len(), - ); + let views = Args::views_no_constants(&columns); + verify_lengths::(&columns, views.as_ref(), row_count)?; - let views = Args::per_row_views(&columns); - ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + let constants = Args::constants(&columns); let prepared = prepare(constants); + // 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 every slot before skipping rows. Recheck addressability afterward because the - // initializer mutably borrows the row representation. 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 = >::row_count(&rows); vortex_ensure_eq!( initialized_row_count, @@ -156,47 +161,100 @@ where "the initialized output sink must address exactly {row_count} rows, got {initialized_row_count}", ); - // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first - // immediate error, turn later callbacks into no-ops, and return before finishing the sink. - let mut error = None; - valid.for_each_set_index(|index| { - if error.is_some() { - return; - } + if let Some(views) = views { + 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 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) }; - let result = match &views { - Some(views) => { - // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows, and - // mask indices are below `row_count`. - let elements = unsafe { Args::get_from_views_unchecked(views, index) }; - apply(&prepared, elements, output) - } - None => apply(&prepared, Args::get(&columns, index), output), - }; - if let Err(row_error) = result.into_result() { - error = Some(row_error); - } - }); + // SAFETY: `verify_lengths` proved every view has `row_count` rows, and mask indices + // are below `row_count`. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; - if let Some(error) = error { - return Err(error); + apply(&prepared, elements, output).into_result() + })?; + } else { + 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() + })?; } } - finish_sink::(sink).map(Some) + // 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) } -fn finish_sink(sink: Sink) -> VortexResult +/// State resolved before preparing the skip-invalid row loop. +struct ValidRowsSetup<'valid, Args, Sink, Options> where + Args: ElementTuple, Sink: OutputSink, { - // SAFETY: callers reach this helper only after every completed callback returned the sink's - // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. - // The sink contract defines how that evidence establishes initialization of its row storage. - unsafe { >::finish(sink) }.map(RowExecution::Output) + 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)] @@ -313,6 +371,7 @@ mod tests { )?; assert!(execution.is_none()); + Ok(()) } @@ -345,6 +404,7 @@ mod tests { let expected = PrimitiveArray::from_iter([20_i64, 0, 60]); assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) } @@ -375,6 +435,7 @@ mod tests { .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/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index e6f63628e88..7c1baed168a 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 @@ -138,7 +138,7 @@ 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. + /// Borrowed views of decoded columns with no batch constants. type Views<'a>; /// The borrowed row of element values. @@ -146,9 +146,9 @@ pub trait ElementTuple: 'static + private::Sealed { /// 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 a non-constant argument. + /// 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>; @@ -186,22 +186,25 @@ pub trait ElementTuple: 'static + private::Sealed { ) -> 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 when none 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 column is batch-constant. Otherwise, omitting [`ArgColumn`] from the + /// returned tuple removes constant checks from the row loop. + fn views_no_constants(columns: &Self::Columns) -> Option>; /// Whether every view 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. + /// The executor calls this once before the loop used when no input is batch-constant. 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. + /// Whether every non-constant argument contains exactly `row_count` rows. /// /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch @@ -266,7 +269,7 @@ impl ElementTuple for () { fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} - fn per_row_views(_columns: &Self::Columns) -> Option> { + fn views_no_constants(_columns: &Self::Columns) -> Option> { Some(()) } @@ -350,7 +353,7 @@ macro_rules! element_tuple { ($(columns.$idx.get(index),)+) } - fn per_row_views(columns: &Self::Columns) -> Option> { + fn views_no_constants(columns: &Self::Columns) -> Option> { Some(($($t::view(columns.$idx.per_row_column()?),)+)) } 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..832dbf9ee19 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 @@ -17,7 +17,7 @@ use crate::scalar_fn::unstable::row::InputElement; /// 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 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 From 8e5e17130daff742445552a040ec68765b354a6b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 17 Aug 2026 15:05:47 -0400 Subject: [PATCH 3/5] address review comments Signed-off-by: Connor Tsui --- .../scalar_fn/unstable/row/execute/owned.rs | 4 +- .../scalar_fn/unstable/row/execute/sink.rs | 91 +++++++------- .../src/scalar_fn/unstable/row/mod.rs | 3 + .../unstable/row/types/element/bool.rs | 8 +- .../unstable/row/types/element/input.rs | 19 ++- .../unstable/row/types/element/mod.rs | 1 - .../unstable/row/types/element/output.rs | 2 +- .../unstable/row/types/element/primitive.rs | 8 +- .../row/types/element/tuple/element_tuple.rs | 118 ++++++++++-------- .../row/types/element/tuple/indexed.rs | 3 +- .../unstable/row/types/element/tuple/mod.rs | 1 - .../unstable/row/types/element/tuple/tests.rs | 26 ++-- .../src/scalar_fn/unstable/row/types/mod.rs | 11 +- .../src/scalar_fn/unstable/row/types/sink.rs | 28 ++--- .../src/scalar_fn/unstable/row/types/view.rs | 56 +++++++++ .../scalar_fn/unstable/row/visitor/check.rs | 4 +- .../src/scalar_fn/unstable/row/visitor/mod.rs | 5 +- .../scalar_fn/unstable/row/visitor/plan.rs | 8 +- .../unstable/row/visitor/row_visitor.rs | 11 +- 19 files changed, 226 insertions(+), 181 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/types/view.rs 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 d3809f939be..1b902f9ecdf 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -67,13 +67,13 @@ where const { assert_owned_output_needs_no_drop::() }; let columns = Args::decode(args, ctx)?; - let prepared = prepare(Args::constants(&columns)); + 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_no_constants(&columns) { + 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::view_lens_match(&views, row_count), 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 4fd2b42476b..2509240508f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -10,7 +10,6 @@ use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -21,29 +20,7 @@ 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; - -/// Verify that every decoded input addresses exactly `row_count` rows. -/// -/// Unlike the owned executor, the paths with and without batch constants can share this check -/// without losing sink-loop vectorization under multiple CGUs without LTO. The exact pass -/// interaction is unknown. -fn verify_lengths( - columns: &Args::Columns, - views: Option<&Args::Views<'_>>, - row_count: usize, -) -> VortexResult<()> { - let lengths_match = match views { - Some(views) => Args::view_lens_match(views, row_count), - None => Args::decoded_lens_match(columns, row_count), - }; - - vortex_ensure!( - lengths_match, - "a decoded row input does not address exactly {row_count} rows", - ); - - Ok(()) -} +use crate::scalar_fn::unstable::row::ViewLen; /// Decode inputs once, then write one sink row for each input row. /// @@ -62,13 +39,11 @@ where ApplyResult: SinkResult>::WriteToken>, { let columns = Args::decode(args, ctx)?; - let views = Args::views_no_constants(&columns); + let views = Args::views_if_no_consts(&columns); let row_count = args.row_count(); - verify_lengths::(&columns, views.as_ref(), row_count)?; - - let constants = Args::constants(&columns); - let prepared = prepare(constants); + let const_values = Args::const_values(&columns); + let prepared = prepare(const_values); let mut sink = >::with_capacity(row_count)?; @@ -77,7 +52,7 @@ where let mut rows = >::rows(&mut sink); // This equality proves to LLVM that `0..row_count` is in bounds for `rows`. - let sink_row_count = >::row_count(&rows); + let sink_row_count = rows.len(); vortex_ensure_eq!( sink_row_count, row_count, @@ -85,8 +60,12 @@ where ); if let Some(views) = views { + if !Args::view_lens_match(&views, row_count) { + decoded_length_error(row_count)?; + } + for index in 0..row_count { - // SAFETY: `verify_lengths` proved every view has `row_count` rows before the loop. + // SAFETY: `view_lens_match` proved every view has `row_count` rows before the loop. 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 = @@ -95,6 +74,10 @@ where 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 = @@ -139,11 +122,9 @@ where return Ok(None); }; - let views = Args::views_no_constants(&columns); - verify_lengths::(&columns, views.as_ref(), row_count)?; - - let constants = Args::constants(&columns); - let prepared = prepare(constants); + 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. @@ -154,7 +135,7 @@ where // The initializer can change addressability. Recheck it so LLVM can prove every mask // index is in bounds. - let initialized_row_count = >::row_count(&rows); + let initialized_row_count = rows.len(); vortex_ensure_eq!( initialized_row_count, row_count, @@ -162,19 +143,27 @@ where ); if let Some(views) = views { + if !Args::view_lens_match(&views, 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: `verify_lengths` proved every view has `row_count` rows, and mask indices - // are below `row_count`. + // SAFETY: `view_lens_match` proved every view has `row_count` rows, and mask + // indices are below `row_count`. 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`. @@ -193,6 +182,18 @@ where .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 @@ -292,7 +293,7 @@ mod tests { type Row<'a> = (); type WriteToken = (); - fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn return_dtype(_options: &Options) -> VortexResult { Ok(DType::from(i64::PTYPE)) } @@ -304,10 +305,6 @@ mod tests { fn rows(&mut self) -> Self::Rows<'_> {} - fn row_count(_rows: &Self::Rows<'_>) -> usize { - 0 - } - unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { } @@ -330,7 +327,7 @@ mod tests { }) } - fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn return_dtype(_options: &Options) -> VortexResult { Ok(DType::from(i64::PTYPE)) } @@ -342,10 +339,6 @@ mod tests { &mut self.0 } - fn row_count(rows: &Self::Rows<'_>) -> usize { - rows.len() - } - unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { &mut rows[index] } diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 057b5cda694..433a20babc4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -17,6 +17,8 @@ //! 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; @@ -32,6 +34,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/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 c9140c00b8c..a1f69dcc454 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,13 +11,14 @@ 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 +/// For every view returned by [`view`](Self::view), every index below [`ViewLen::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. pub unsafe trait InputElement: 'static { @@ -29,7 +30,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 +42,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<()>; @@ -89,12 +90,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 +99,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/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 7a8a8e92e41..51d66594332 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -20,4 +20,3 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; -pub use tuple::batch_constant; 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 7c1baed168a..2b576413473 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,62 +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 { - /// One decoded value per batch row; executors validate the exact length before traversal. - PerRow(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_constant`]. - Constant(T::Column), + /// Exactly one decoded row, established by [`ArgColumn::try_from_const`]. + Const(T::Column), } impl ArgColumn { - fn try_from_constant(column: T::Column) -> VortexResult { - let decoded_len = T::view_len(&T::view(&column)); + 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::Constant(column))) + 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 Self::try_from_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 Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?).map(Some); + 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); } @@ -83,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 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_value(&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)), } } } @@ -115,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()); } - array - .as_opt::() - .is_some_and(|ext| ext.storage_array().is::()) - .then(|| array.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()); + } + + None } /// Typed argument tuples for arities zero through twelve. @@ -146,9 +155,9 @@ pub trait ElementTuple: 'static + private::Sealed { /// The batch-constant element values. /// - /// `Some` carries the value of a batch-constant argument. `None` marks a non-constant 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>; @@ -159,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<()>; @@ -191,11 +200,11 @@ pub trait ElementTuple: 'static + private::Sealed { /// visible in the loop so LLVM can unswitch it before vectorizing. fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; - /// Borrow the decoded columns when none is batch-constant. + /// Borrow the decoded columns only when no argument is batch-constant. /// - /// Returns `None` if any column is batch-constant. Otherwise, omitting [`ArgColumn`] from the - /// returned tuple removes constant checks from the row loop. - fn views_no_constants(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. /// @@ -204,7 +213,7 @@ pub trait ElementTuple: 'static + private::Sealed { /// 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 non-constant argument contains exactly `row_count` rows. + /// Whether every argument decoded at full batch length contains exactly `row_count` rows. /// /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch @@ -225,9 +234,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 () {} @@ -240,7 +252,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!( @@ -269,7 +281,7 @@ impl ElementTuple for () { fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} - fn views_no_constants(_columns: &Self::Columns) -> Option> { + fn views_if_no_consts(_columns: &Self::Columns) -> Option> { Some(()) } @@ -289,7 +301,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 { @@ -304,7 +316,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!( @@ -353,15 +365,15 @@ macro_rules! element_tuple { ($(columns.$idx.get(index),)+) } - fn views_no_constants(columns: &Self::Columns) -> Option> { - Some(($($t::view(columns.$idx.per_row_column()?),)+)) + fn views_if_no_consts(columns: &Self::Columns) -> Option> { + Some(($($t::view(columns.$idx.as_column()?),)+)) } fn view_lens_match( views: &Self::Views<'_>, row_count: usize, ) -> bool { - $($t::view_len(&views.$idx) == row_count &&)+ true + $(views.$idx.len() == row_count &&)+ true } fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { @@ -383,8 +395,8 @@ macro_rules! element_tuple { ($(unsafe { $t::get_from_view_unchecked(&views.$idx, index) },)+) } - fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { - ($(columns.$idx.constant_value(),)+) + 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 832dbf9ee19..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,6 +11,7 @@ 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. /// @@ -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/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index 69b5cf686f6..a2c143704a0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -8,7 +8,6 @@ mod element_tuple; pub use element_tuple::ElementTuple; -pub use element_tuple::batch_constant; mod indexed; pub use indexed::IndexedElementTuple; 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 560e1e48f4f..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::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 e47f195410c..b1e419da195 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -3,16 +3,16 @@ //! 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 the immediate and deferred +//! outcomes returned by sink-writing row closures. mod element; pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; -pub(super) use element::batch_constant; mod result; pub use result::SinkResult; @@ -21,3 +21,6 @@ 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/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index 22c90b8da02..dd702fd5cd8 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -14,11 +14,12 @@ 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. /// Execution can omit invalid rows when [`skipped_rows_initializer`] returns an initializer. @@ -33,7 +34,8 @@ 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 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. @@ -47,7 +49,6 @@ use crate::scalar_fn::unstable::row::OutputElement; /// [`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 @@ -56,7 +57,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; @@ -83,11 +84,11 @@ pub unsafe trait OutputSink: 'static + Sized { 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 +96,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 /// @@ -183,7 +181,7 @@ unsafe impl OutputSink }) } - fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn return_dtype(_options: &Options) -> VortexResult { Ok(T::element_dtype()) } @@ -198,10 +196,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..d9063fdcde7 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/view.rs @@ -0,0 +1,56 @@ +// 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) + } +} 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..40702d68a03 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -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", ); } @@ -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/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index c7f9baf6a62..1b063c788db 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -8,12 +8,11 @@ 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; -pub(super) use execute::ExecuteRows; -pub(super) use execute::ExecuteValidRows; mod plan; -pub(super) use plan::BatchPlan; pub(super) use plan::BatchPlanner; pub(super) use plan::RowPolicy; 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 5360ced573b..d41c7cd7ce6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -114,6 +114,8 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. + // TODO(connor)[RowFn]: Remove this expectation when #9450 uses the planned policy. + #[expect(dead_code)] pub(crate) policy: RowPolicy, } @@ -143,7 +145,7 @@ pub(crate) enum RowPolicy { 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 @@ -152,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 @@ -161,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 1a9e8628bc7..ebf56594fb5 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 @@ -74,7 +74,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 per-row 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 +162,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) } From 6f16c60f39155a25cc4bf98c2c0a33839667209b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 17 Aug 2026 16:42:21 -0400 Subject: [PATCH 4/5] Address remaining RowFn review comments Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/execute/owned.rs | 3 ++- vortex-array/src/scalar_fn/unstable/row/mod.rs | 1 + .../unstable/row/types/element/input.rs | 10 ++++++---- .../row/types/element/tuple/element_tuple.rs | 4 ++-- .../src/scalar_fn/unstable/row/types/mod.rs | 5 +++-- .../src/scalar_fn/unstable/row/types/result.rs | 17 ++++++++++++++--- .../src/scalar_fn/unstable/row/types/sink.rs | 7 ++++++- .../src/scalar_fn/unstable/row/visitor/check.rs | 4 ++-- .../scalar_fn/unstable/row/visitor/execute.rs | 7 +++---- .../src/scalar_fn/unstable/row/visitor/plan.rs | 4 ++-- .../unstable/row/visitor/row_visitor.rs | 12 +++++------- 11 files changed, 46 insertions(+), 28 deletions(-) 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 1b902f9ecdf..44392de307b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -16,6 +16,7 @@ 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::visitor::assert_owned_output_needs_no_drop; @@ -59,7 +60,7 @@ pub(crate) fn execute_owned( where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + 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 diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 433a20babc4..12ce0db87c2 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -27,6 +27,7 @@ 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; 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 a1f69dcc454..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 @@ -17,10 +17,12 @@ use crate::scalar_fn::unstable::row::ViewLen; /// /// # Safety /// -/// For every view returned by [`view`](Self::view), every index below [`ViewLen::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; 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 2b576413473..96aad70f24e 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 @@ -202,7 +202,7 @@ pub trait ElementTuple: 'static + private::Sealed { /// Borrow the decoded columns only when no argument is batch-constant. /// - /// Returns `None` if any argument is batch-constant. Otherwise, omitting [`ArgColumn`] from the + /// 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>; @@ -217,7 +217,7 @@ pub trait ElementTuple: 'static + private::Sealed { /// /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch - /// constant is exempt because its [`ArgColumn`] constructor already validated the one-row + /// 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; 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 b1e419da195..3b4a92e784e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -5,8 +5,8 @@ //! //! [`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 the immediate and deferred -//! outcomes returned by sink-writing row closures. +//! that need row handles or batch-wide state. [`result`] defines immediate and deferred row +//! outcomes. mod element; pub use element::ElementTuple; @@ -15,6 +15,7 @@ pub use element::InputElement; pub use element::OutputElement; mod result; +pub use result::FailureEvidence; pub use result::SinkResult; mod sink; 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 dd702fd5cd8..e95431bf9e6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -36,15 +36,20 @@ use crate::scalar_fn::unstable::row::ViewLen; /// /// - 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 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 40702d68a03..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; @@ -69,7 +69,7 @@ where Function: RowFn, Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + Fail: FailureEvidence, { assert_owned_visit_contract::(); assert!( 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 b003f70a395..610568343f3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -7,8 +7,6 @@ //! match the plan before entering a row loop. [`ExecuteValidRows`] can decline, so the batch layer //! filters the inputs and retries with [`ExecuteRows`]. -use std::ops::BitOrAssign; - use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; @@ -25,6 +23,7 @@ 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; @@ -138,7 +137,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + Fail: FailureEvidence, { const { assert_deferred_visit_contract::() }; ensure_plan( @@ -270,7 +269,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + Fail: FailureEvidence, { const { assert_deferred_visit_contract::() }; ensure_plan( 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 d41c7cd7ce6..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 { 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 ebf56594fb5..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; @@ -197,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 @@ -238,7 +236,7 @@ pub trait RowVisitor: private::Sealed + Sized { where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign, + Fail: FailureEvidence, { self.visit_prepared_deferred::( |_| (), @@ -291,7 +289,7 @@ pub trait RowVisitor: private::Sealed + Sized { where Args: IndexedElementTuple, Out: OutputElement, - Fail: Copy + Default + BitOrAssign; + Fail: FailureEvidence; } pub(super) mod private { From d65a9da7b532e795672ae44e223c3280f15d01c3 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 18 Aug 2026 11:03:21 -0400 Subject: [PATCH 5/5] final touches Signed-off-by: Connor Tsui --- .../scalar_fn/unstable/row/execute/owned.rs | 9 ++-- .../scalar_fn/unstable/row/execute/sink.rs | 14 +++--- .../row/types/element/tuple/element_tuple.rs | 31 ++++-------- .../src/scalar_fn/unstable/row/types/view.rs | 48 +++++++++++++++++++ 4 files changed, 69 insertions(+), 33 deletions(-) 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 44392de307b..513178cd97c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -19,6 +19,7 @@ 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. @@ -77,19 +78,19 @@ where 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::view_lens_match(&views, row_count), + Args::ARITY == 0 || views.len() == row_count, "a decoded row input does not address exactly {row_count} rows", ); - // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows - // immediately above. + // 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. The exact pass interaction is unknown. + // LTO. vortex_ensure!( Args::decoded_lens_match(&columns, row_count), "a decoded row input does not address exactly {row_count} rows", 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 2509240508f..9b809179cae 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -39,7 +39,6 @@ where ApplyResult: SinkResult>::WriteToken>, { let columns = Args::decode(args, ctx)?; - let views = Args::views_if_no_consts(&columns); let row_count = args.row_count(); let const_values = Args::const_values(&columns); @@ -59,13 +58,15 @@ where "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::view_lens_match(&views, row_count) { + if Args::ARITY != 0 && views.len() != row_count { decoded_length_error(row_count)?; } for index in 0..row_count { - // SAFETY: `view_lens_match` proved every view has `row_count` rows before the loop. + // 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 = @@ -143,7 +144,7 @@ where ); if let Some(views) = views { - if !Args::view_lens_match(&views, row_count) { + if Args::ARITY != 0 && views.len() != row_count { decoded_length_error(row_count)?; } @@ -153,8 +154,9 @@ where let output = unsafe { >::row_unchecked(&mut rows, index) }; - // SAFETY: `view_lens_match` proved every view has `row_count` rows, and mask - // indices are below `row_count`. + // 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() 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 96aad70f24e..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 @@ -148,7 +148,7 @@ pub trait ElementTuple: 'static + private::Sealed { type Columns; /// Borrowed views of decoded columns with no batch constants. - type Views<'a>; + type Views<'a>: ViewLen; /// The borrowed row of element values. type Elems<'a>; @@ -206,19 +206,15 @@ pub trait ElementTuple: 'static + private::Sealed { /// 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. - /// - /// The executor calls this once before the loop used when no input is batch-constant. 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 argument decoded at full batch length contains exactly `row_count` rows. /// - /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include - /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch - /// constant is exempt because its `ArgColumn` constructor already validated the one-row - /// representation produced by decoding. + /// 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 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. @@ -285,10 +281,6 @@ impl ElementTuple for () { 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 } @@ -369,13 +361,6 @@ macro_rules! element_tuple { Some(($($t::view(columns.$idx.as_column()?),)+)) } - fn view_lens_match( - views: &Self::Views<'_>, - row_count: usize, - ) -> bool { - $(views.$idx.len() == row_count &&)+ true - } - fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { $(columns.$idx.addresses_rows(row_count) &&)+ true } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/view.rs b/vortex-array/src/scalar_fn/unstable/row/types/view.rs index d9063fdcde7..18085a12d14 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/view.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/view.rs @@ -54,3 +54,51 @@ impl ViewLen for &mut T { 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(); + } +}