diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs index 91700a79d23..ff7fbe7ab81 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -116,7 +116,7 @@ where // vectorization. Check each divide immediately and stop at the first failure. // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when // they need to skip invalid rows. - visitor.visit_into::<(T, T), UninitElementSink, _>(|(lhs, rhs), output| { + visitor.visit_into::<(T, T), UninitElementSink, _>((), |(lhs, rhs), output| { let (value, failed) = CheckedDiv::apply(lhs, rhs); if failed { return Err(numeric_error(>::ERROR)); diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 224c875c52b..25a260846c6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -9,6 +9,8 @@ use rstest::rstest; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_session::registry::CachedId; use super::finalize_kernel_output; @@ -18,6 +20,7 @@ use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -30,6 +33,8 @@ use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::FixedSizeListSink; +use crate::scalar_fn::unstable::row::InitializedRow; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; @@ -78,15 +83,16 @@ struct I64Sink(BufferMut); // SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that // initialized slice. The `()` write token therefore proves no additional invariant. unsafe impl OutputSink for I64Sink { + type Params = (); type Rows<'a> = &'a mut [i64]; type Row<'a> = &'a mut i64; type WriteToken = (); - fn storage_dtype() -> DType { + fn storage_dtype(_params: &Self::Params) -> DType { DType::from(i64::PTYPE) } - fn with_capacity(rows: usize) -> VortexResult { + fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { Ok(Self(BufferMut::zeroed(rows))) } @@ -104,6 +110,38 @@ unsafe impl OutputSink for I64Sink { } } +#[derive(Clone)] +struct RepeatValue; + +impl RowFn for RepeatValue { + type Options = usize; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.repeat_value"); + *ID + } + + fn dispatch( + &self, + width: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + vortex_ensure!( + u32::try_from(*width).is_ok(), + InvalidArgument: + "test.repeat_value width must fit in the fixed-size-list u32 list size, got {width}", + ); + + visitor.visit_into::<(i64,), FixedSizeListSink, _>(*width, |(value,), row| { + InitializedRow::fill(row, |_| value) + }) + } +} + impl RowFn for DeferredAdd { type Options = EmptyOptions; @@ -156,7 +194,7 @@ impl RowFn for ValidOnlyIdentity { _args: &[DType], visitor: V, ) -> VortexResult { - visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| { + visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>((), |(value,), output| { *output = value; Ok(()) }) @@ -184,6 +222,58 @@ impl RowFn for InvalidKernelOutput { } } +#[rstest] +#[case::dense_width_two( + vec![1_i64, 2], + Validity::NonNullable, + 2, + vec![1_i64, 1, 2, 2], +)] +#[case::dense_width_four( + vec![3_i64, 4], + Validity::NonNullable, + 4, + vec![3_i64, 3, 3, 3, 4, 4, 4, 4], +)] +#[case::empty(vec![], Validity::NonNullable, 3, vec![])] +#[case::zero_width(vec![3_i64, 4], Validity::NonNullable, 0, vec![])] +#[case::all_null( + vec![5_i64, 6], + Validity::AllInvalid, + 2, + vec![0_i64, 0, 0, 0], +)] +#[case::partially_valid( + vec![7_i64, 8, 9], + Validity::from_iter([true, false, true]), + 3, + vec![7_i64, 7, 7, 0, 0, 0, 9, 9, 9], +)] +fn test_fixed_size_list_sink_uses_runtime_width( + #[case] input_values: Vec, + #[case] validity: Validity, + #[case] width: usize, + #[case] expected_elements: Vec, +) -> VortexResult<()> { + let row_count = input_values.len(); + let input = PrimitiveArray::new(input_values, validity.clone()).into_array(); + let args = VecExecutionArgs::new(vec![input], row_count); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&RepeatValue, &width, &args, &mut ctx)?; + let expected = FixedSizeListArray::new( + PrimitiveArray::from_iter(expected_elements).into_array(), + u32::try_from(width) + .map_err(|_| vortex_err!(InvalidArgument: "test width must fit in u32, got {width}"))?, + validity, + row_count, + ) + .into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + #[test] fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<()> { static ID: CachedId = CachedId::new("test.finalize_kernel_output"); @@ -403,7 +493,7 @@ impl RowFn for DeclaredSinkOutput { ) -> VortexResult { visitor .with_output_dtype(timestamp_dtype(Nullability::NonNullable)) - .visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| { + .visit_into::<(i64,), I64Sink, VortexResult<()>>((), |(value,), output| { *output = value; Ok(()) }) 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 7ed012191ea..9b9d3b196f1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -28,6 +28,7 @@ use crate::scalar_fn::unstable::row::ViewLen; /// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant. pub(crate) fn execute_sink( args: &dyn ExecutionArgs, + params: &Sink::Params, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, @@ -43,7 +44,7 @@ where let const_values = Args::const_values(&columns); let prepared = prepare(const_values); - let mut sink = Sink::with_capacity(row_count)?; + let mut sink = Sink::with_capacity(row_count, params)?; // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. { @@ -100,6 +101,7 @@ where pub(crate) fn execute_sink_valid_rows( args: &dyn ExecutionArgs, valid: &MaskValuesRef, + params: &Sink::Params, ctx: &mut ExecutionCtx, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, @@ -115,7 +117,7 @@ where valid_rows, row_count, mut sink, - }) = setup_sink_valid_rows::(args, valid, ctx)? + }) = setup_sink_valid_rows::(args, valid, params, ctx)? else { return Ok(None); }; @@ -205,6 +207,7 @@ where fn setup_sink_valid_rows<'valid, Args, Sink>( args: &dyn ExecutionArgs, valid: &'valid MaskValuesRef, + params: &Sink::Params, ctx: &mut ExecutionCtx, ) -> VortexResult>> where @@ -227,7 +230,7 @@ where // 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 = Sink::with_capacity(row_count)?; + let sink = Sink::with_capacity(row_count, params)?; let valid_rows = valid.bit_buffer(); vortex_ensure_eq!( @@ -272,15 +275,16 @@ mod tests { // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or // `finish` through the executor. The row-initialization requirements are therefore vacuous. unsafe impl OutputSink for NonSkippingSink { + type Params = (); type Rows<'a> = (); type Row<'a> = (); type WriteToken = (); - fn storage_dtype() -> DType { + fn storage_dtype(_params: &Self::Params) -> DType { DType::from(i64::PTYPE) } - fn with_capacity(_rows: usize) -> VortexResult { + fn with_capacity(_rows: usize, _params: &Self::Params) -> VortexResult { Err(vortex_err!( "a non-skipping sink must decline before allocation" )) @@ -300,6 +304,7 @@ mod tests { // 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 Params = (); type Rows<'a> = &'a mut Vec; type Row<'a> = &'a mut i64; type WriteToken = (); @@ -310,11 +315,11 @@ mod tests { }) } - fn storage_dtype() -> DType { + fn storage_dtype(_params: &Self::Params) -> DType { DType::from(i64::PTYPE) } - fn with_capacity(rows: usize) -> VortexResult { + fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { Ok(Self(vec![0; rows])) } @@ -343,6 +348,7 @@ mod tests { let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>( &args, &valid, + &(), &mut ctx, |_| (), |_, _, _| (), @@ -365,6 +371,7 @@ mod tests { let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, ()>( &args, &valid, + &(), &mut ctx, |_| (), |_, (value,), output| { diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index bd73916c85c..8bf45929163 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -31,8 +31,10 @@ pub use row_fn::RowFn; mod types; pub use types::ElementTuple; pub use types::FailureEvidence; +pub use types::FixedSizeListSink; pub use types::IndexedElementTuple; pub use types::InitializedElement; +pub use types::InitializedRow; pub use types::InputElement; pub use types::OutputElement; pub use types::OutputSink; 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 4fc0c323ad4..007d9916d26 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -20,7 +20,9 @@ pub use result::FailureEvidence; pub use result::SinkResult; mod sink; +pub use sink::FixedSizeListSink; pub use sink::InitializedElement; +pub use sink::InitializedRow; pub use sink::OutputSink; pub use sink::UninitElementSink; 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 42914150691..469b6443ce0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/result.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/result.rs @@ -4,14 +4,15 @@ //! Return types for row closures. //! //! [`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. +//! the executor handle initialized sinks and sinks that require an [`InitializedElement`] or +//! [`InitializedRow`] token, with either infallible or immediate-error callbacks. use std::ops::BitOrAssign; use vortex_error::VortexResult; use super::InitializedElement; +use super::InitializedRow; /// Compact failure evidence that can be OR-reduced across rows. /// @@ -57,6 +58,17 @@ impl SinkResult for InitializedElement { } } +impl private::Sealed for InitializedRow {} + +impl SinkResult for InitializedRow { + type WriteToken = InitializedRow; + const INFALLIBLE: bool = true; + + fn into_result(self) -> VortexResult<()> { + Ok(()) + } +} + impl private::Sealed for VortexResult<()> {} impl SinkResult for VortexResult<()> { diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs new file mode 100644 index 00000000000..6754983e2f8 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Runtime-width fixed-size-list output for row kernels. +//! +//! [`FixedSizeListSink`] stores each row in a contiguous slice of one flat element allocation. +//! [`InitializedRow`] proves that a row callback filled its entire slice before finishing. + +use std::mem::MaybeUninit; +use std::sync::Arc; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::OutputSink; +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::FixedSizeListArray; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::ViewLen; +use crate::validity::Validity; + +/// Proof that every element in one uninitialized fixed-size row was initialized. +/// +/// The private field prevents construction without calling [`fill`](Self::fill), which writes the +/// complete row before returning. +#[must_use = "return this token from the row closure to prove that it initialized the output"] +pub struct InitializedRow(()); + +impl InitializedRow { + /// Fill every element in `row` and return its proof token. + #[inline] + pub fn fill( + row: &mut [MaybeUninit], + mut value_for_index: impl FnMut(usize) -> T, + ) -> Self { + for (index, element) in row.iter_mut().enumerate() { + element.write(value_for_index(index)); + } + + Self(()) + } +} + +/// A loop-local view of the flat storage and runtime shape of a [`FixedSizeListSink`]. +pub struct FixedSizeRows<'a, T> { + /// The flat elements for all output rows. + elements: &'a mut [MaybeUninit], + + /// The number of elements in each output row. + width: usize, + + /// The number of output rows, stored separately so zero-width rows remain addressable. + row_count: usize, +} + +impl ViewLen for FixedSizeRows<'_, T> { + fn len(&self) -> usize { + self.row_count + } +} + +/// A fixed-size-list sink whose row width is supplied at dispatch time. +/// +/// The row closure must return the [`InitializedRow`] from [`InitializedRow::fill`] on success. +/// The width must fit in the `u32` list size stored by [`FixedSizeListArray`]. A dispatch derives +/// and validates that physical parameter before calling [`RowVisitor::visit_into`]. +/// +/// [`RowVisitor::visit_into`]: crate::scalar_fn::unstable::row::RowVisitor::visit_into +pub struct FixedSizeListSink { + /// Spare flat storage written one fixed-size row at a time. + values: Vec, + + /// The number of elements in each output row. + width: usize, + + /// The number of output rows. + row_count: usize, +} + +// SAFETY: `with_capacity` reserves `row_count * width` elements, and `FixedSizeRows` retains that +// shape for its lifetime. Each row is one disjoint `width`-element slice. `InitializedRow::fill` +// writes every element before returning its private token, and the skipped-row initializer writes +// every flat element before masked traversal. `values` retains length zero until every row is safe +// to publish in `finish`. +unsafe impl OutputSink for FixedSizeListSink { + type Params = usize; + type Rows<'a> = FixedSizeRows<'a, T>; + type Row<'a> = &'a mut [MaybeUninit]; + type WriteToken = InitializedRow; + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + for element in rows.elements.iter_mut() { + element.write(T::default()); + } + }) + } + + fn storage_dtype(params: &Self::Params) -> DType { + DType::FixedSizeList( + Arc::new(T::element_dtype()), + fixed_size_list_size(*params), + Nullability::NonNullable, + ) + } + + fn with_capacity(rows: usize, params: &Self::Params) -> VortexResult { + let width = *params; + let element_capacity = rows.checked_mul(width).ok_or_else(|| { + vortex_err!( + InvalidArgument: + "fixed-size-list sink capacity must fit in usize, got {rows} rows with width {width}" + ) + })?; + + Ok(Self { + values: Vec::with_capacity(element_capacity), + width, + row_count: rows, + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + FixedSizeRows { + elements: &mut self.values.spare_capacity_mut()[..self.row_count * self.width], + width: self.width, + row_count: self.row_count, + } + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + let start = index * rows.width; + let end = start + rows.width; + + // SAFETY: the caller guarantees `index < row_count`, and construction guarantees the + // element slice has length `row_count * width`. + unsafe { rows.elements.get_unchecked_mut(start..end) } + } + + unsafe fn finish(mut self) -> VortexResult { + let element_count = self.row_count * self.width; + + // SAFETY: the caller guarantees every row was initialized, and `with_capacity` reserved + // `row_count * width` elements. + unsafe { self.values.set_len(element_count) }; + + let elements = T::build(self.values); + let lists = FixedSizeListArray::new( + elements, + fixed_size_list_size(self.width), + Validity::NonNullable, + self.row_count, + ); + + Ok(lists.into_array()) + } +} + +fn fixed_size_list_size(width: usize) -> u32 { + u32::try_from(width) + .vortex_expect("fixed-size-list sink width must fit in u32; dispatch validated it") +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs similarity index 58% rename from vortex-array/src/scalar_fn/unstable/row/types/sink.rs rename to vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs index 40f449ff4ae..12cc1db7e5d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs @@ -3,19 +3,23 @@ //! Output builders for row kernels that cannot return independent owned values. //! -//! [`OutputSink`] allocates batch-wide state and lends one row handle to each callback. -//! [`UninitElementSink`] is the fixed-width implementation used when avoiding output -//! initialization matters. - -use std::mem::MaybeUninit; +//! [`OutputSink`] owns the shared lifecycle and safety contract. [`UninitElementSink`] provides +//! uninitialized scalar storage, while [`FixedSizeListSink`] provides runtime-width row storage. 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; +mod fixed_size_list; +pub use fixed_size_list::FixedSizeListSink; +pub use fixed_size_list::InitializedRow; + +mod uninit_element; +pub use uninit_element::InitializedElement; +pub use uninit_element::UninitElementSink; + /// A column allocated once per batch that a row closure writes into, one row at a time. /// /// A sink owns batch-wide state that an independent owned value cannot express, such as @@ -64,6 +68,12 @@ use crate::scalar_fn::unstable::row::ViewLen; /// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult /// [`skipped_rows_initializer`]: Self::skipped_rows_initializer pub unsafe trait OutputSink: 'static + Sized { + /// Physical parameters required to construct this sink before the row loop. + /// + /// This type describes only physical storage. A logical output dtype belongs on + /// [`RowVisitor::with_output_dtype`](crate::scalar_fn::unstable::row::RowVisitor::with_output_dtype). + type Params: 'static; + /// A loop-local view of all output rows. /// /// Borrowed once before execution so the sink's buffer descriptor and shape become loop @@ -97,16 +107,12 @@ pub unsafe trait OutputSink: 'static + Sized { /// The dtype of the column this sink builds. /// - /// Because this method takes no arguments, the dtype must be a property of the Rust type. An - /// output dtype that depends on the function options or argument dtypes is declared by - /// [`RowVisitor::with_output_dtype`](crate::scalar_fn::unstable::row::RowVisitor::with_output_dtype). - /// /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the /// result, and masks the null rows. - fn storage_dtype() -> DType; + fn storage_dtype(params: &Self::Params) -> DType; /// Allocate a sink for `rows` rows. - fn with_capacity(rows: usize) -> VortexResult; + fn with_capacity(rows: usize, params: &Self::Params) -> VortexResult; /// Borrow all output rows for the hot loop. fn rows(&mut self) -> Self::Rows<'_>; @@ -119,7 +125,8 @@ pub unsafe trait OutputSink: 'static + Sized { 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 - /// [`storage_dtype`](Self::storage_dtype). Called once per batch. + /// [`storage_dtype`](Self::storage_dtype) for the parameters passed to + /// [`with_capacity`](Self::with_capacity). Called once per batch. /// /// # Safety /// @@ -129,96 +136,3 @@ pub unsafe trait OutputSink: 'static + Sized { /// [`skipped_rows_initializer`](Self::skipped_rows_initializer) must have run before traversal. unsafe fn finish(self) -> VortexResult; } - -/// Proof that one uninitialized element row was initialized. -/// -/// The private field prevents safe construction without calling [`write`](Self::write): -/// -/// ```compile_fail,E0423 -/// use vortex_array::scalar_fn::unstable::row::InitializedElement; -/// -/// let _evidence = InitializedElement(()); -/// ``` -#[must_use = "return this token from the row closure to prove that it initialized the output"] -pub struct InitializedElement( - /// Private so constructing initialization evidence requires an unsafe operation. - (), -); - -impl InitializedElement { - /// Write `value` into an uninitialized row and return its proof token. - /// - /// # Safety - /// - /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller - /// must return the token from that callback. Using another row or returning the token from - /// another callback can cause undefined behavior. - #[inline] - pub unsafe fn write(row: &mut MaybeUninit, value: T) -> Self { - row.write(value); - - Self(()) - } -} - -/// An element sink that leaves dense output uninitialized before the row loop. -/// -/// 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. -/// -/// 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 { - /// Spare storage written in increasing row order. - values: Vec, - - /// The number of slots exposed to the row loop and initialized before finishing. - row_count: usize, -} - -// SAFETY: the row slice covers exactly the reserved spare-capacity range, so each accepted index -// names one distinct slot. Safe code cannot construct `InitializedElement`. Its unsafe constructor -// writes the supplied slot and requires the caller to return that exact evidence. The -// skipped-row initializer writes `T::default()` into every slot before masked traversal. -unsafe impl OutputSink for UninitElementSink { - type Rows<'a> = &'a mut [MaybeUninit]; - type Row<'a> = &'a mut MaybeUninit; - type WriteToken = InitializedElement; - - fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { - Some(|rows| { - for row in rows.iter_mut() { - row.write(T::default()); - } - }) - } - - fn storage_dtype() -> DType { - T::element_dtype() - } - - fn with_capacity(rows: usize) -> VortexResult { - Ok(Self { - values: Vec::with_capacity(rows), - row_count: rows, - }) - } - - fn rows(&mut self) -> Self::Rows<'_> { - &mut self.values.spare_capacity_mut()[..self.row_count] - } - - unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { - // SAFETY: required by this method's contract. - unsafe { rows.get_unchecked_mut(index) } - } - - unsafe fn finish(mut self) -> VortexResult { - // SAFETY: the caller guarantees every slot in `0..row_count` was initialized, and - // `with_capacity` reserved every slot in that range. - unsafe { self.values.set_len(self.row_count) }; - - Ok(T::build(self.values)) - } -} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs new file mode 100644 index 00000000000..4087ee94fae --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Uninitialized single-element output for row kernels. +//! +//! [`UninitElementSink`] exposes one uninitialized slot per output row. [`InitializedElement`] +//! proves that a successful callback wrote its slot before the sink publishes the output. + +use std::mem::MaybeUninit; + +use vortex_error::VortexResult; + +use super::OutputSink; +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::OutputElement; + +/// Proof that one uninitialized element row was initialized. +/// +/// The private field prevents safe construction without calling [`write`](Self::write): +/// +/// ```compile_fail,E0423 +/// use vortex_array::scalar_fn::unstable::row::InitializedElement; +/// +/// let _evidence = InitializedElement(()); +/// ``` +#[must_use = "return this token from the row closure to prove that it initialized the output"] +pub struct InitializedElement( + /// Private so constructing initialization evidence requires an unsafe operation. + (), +); + +impl InitializedElement { + /// Write `value` into an uninitialized row and return its proof token. + /// + /// # Safety + /// + /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller + /// must return the token from that callback. Using another row or returning the token from + /// another callback can cause undefined behavior. + #[inline] + pub unsafe fn write(row: &mut MaybeUninit, value: T) -> Self { + row.write(value); + + Self(()) + } +} + +/// An element sink that leaves dense output uninitialized before the row loop. +/// +/// 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. +/// +/// 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 { + /// Spare storage written in increasing row order. + values: Vec, + + /// The number of slots exposed to the row loop and initialized before finishing. + row_count: usize, +} + +// SAFETY: the row slice covers exactly the reserved spare-capacity range, so each accepted index +// names one distinct slot. Safe code cannot construct `InitializedElement`. Its unsafe constructor +// writes the supplied slot and requires the caller to return that exact evidence. The +// skipped-row initializer writes `T::default()` into every slot before masked traversal. +unsafe impl OutputSink for UninitElementSink { + type Params = (); + type Rows<'a> = &'a mut [MaybeUninit]; + type Row<'a> = &'a mut MaybeUninit; + type WriteToken = InitializedElement; + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + for row in rows.iter_mut() { + row.write(T::default()); + } + }) + } + + fn storage_dtype(_params: &Self::Params) -> DType { + T::element_dtype() + } + + fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { + Ok(Self { + values: Vec::with_capacity(rows), + row_count: rows, + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.values.spare_capacity_mut()[..self.row_count] + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(mut self) -> VortexResult { + // SAFETY: the caller guarantees every slot in `0..row_count` was initialized, and + // `with_capacity` reserved every slot in that range. + unsafe { self.values.set_len(self.row_count) }; + + Ok(T::build(self.values)) + } +} 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 cf5bf7ad44f..ecacfddf4c1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -96,14 +96,17 @@ pub(super) fn validate_owned_visit( Ok(dtype) } -pub(super) fn validate_sink_visit(dtypes: &[DType]) -> VortexResult +pub(super) fn validate_sink_visit( + dtypes: &[DType], + params: &Sink::Params, +) -> VortexResult where Args: ElementTuple, Sink: OutputSink, { Args::validate(dtypes)?; - let dtype = Sink::storage_dtype(); + let dtype = Sink::storage_dtype(params); vortex_ensure!( !dtype.is_nullable(), "row output sinks must declare a non-nullable dtype, got {dtype}", diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index 435ae9698d9..0501d6fc6a0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -110,6 +110,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { fn visit_prepared_into( self, + params: Sink::Params, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult @@ -120,13 +121,15 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { { const { assert_sink_visit_contract::() }; let visited = BatchPlan::new( - validate_sink_visit::(self.dtypes)?, + validate_sink_visit::(self.dtypes, ¶ms)?, self.output_dtype, RowPolicy::for_sink::(), )?; self.plan.ensure_reproduced_by(&visited)?; - execute_sink::(self.args, self.ctx, prepare, apply) + execute_sink::( + self.args, ¶ms, self.ctx, prepare, apply, + ) } fn visit_prepared_deferred( @@ -243,6 +246,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { fn visit_prepared_into( self, + params: Sink::Params, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult @@ -253,7 +257,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { { const { assert_sink_visit_contract::() }; let visited = BatchPlan::new( - validate_sink_visit::(self.dtypes)?, + validate_sink_visit::(self.dtypes, ¶ms)?, self.output_dtype, RowPolicy::for_sink::(), )?; @@ -262,6 +266,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { execute_sink_valid_rows::( self.args, &self.valid, + ¶ms, self.ctx, prepare, apply, 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 f603ca6e685..0ebe1a9ad6d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -86,6 +86,7 @@ impl RowVisitor for BatchPlanner<'_, F> { fn visit_prepared_into( self, + params: Sink::Params, _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, _apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult @@ -97,7 +98,7 @@ impl RowVisitor for BatchPlanner<'_, F> { const { assert_sink_visit_contract::() }; BatchPlan::new( - validate_sink_visit::(self.dtypes)?, + validate_sink_visit::(self.dtypes, ¶ms)?, self.output_dtype, RowPolicy::for_sink::(), ) diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs index b4e77a6c50f..93313904b47 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/retry.rs @@ -102,6 +102,7 @@ impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { fn visit_prepared_into( self, + params: Sink::Params, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult @@ -112,14 +113,16 @@ impl RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> { { const { assert_sink_visit_contract::() }; let visited = BatchPlan::new( - validate_sink_visit::(self.args.dtypes())?, + validate_sink_visit::(self.args.dtypes(), ¶ms)?, self.output_dtype, RowPolicy::for_sink::(), )?; self.args.plan().ensure_reproduced_by(&visited)?; - execute_sink::(self.args, self.ctx, prepare, apply) - .map(DenseAttempt::Values) + execute_sink::( + self.args, ¶ms, self.ctx, prepare, apply, + ) + .map(DenseAttempt::Values) } fn visit_prepared_deferred( 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 4de9e34b330..66b3ef03437 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 @@ -170,6 +170,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// } /// /// visitor.visit_into::<(i64, i64), UninitElementSink, _>( + /// (), /// |(lhs, rhs), output| { /// let Some(value) = lhs.checked_div(rhs) else { /// return Err(integer_division_error()); @@ -182,6 +183,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// ``` fn visit_into( self, + params: Sink::Params, apply: impl Fn(Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult where @@ -190,6 +192,7 @@ pub trait RowVisitor: private::Sealed + Sized { ApplyResult: SinkResult, { self.visit_prepared_into::( + params, |_| (), move |&(), args, row| apply(args, row), ) @@ -209,6 +212,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// ConstVectorMagnitudes, /// InitializedElement, /// >( + /// (), /// |(lhs, rhs)| ConstVectorMagnitudes { /// lhs: lhs.map(vector_magnitude), /// rhs: rhs.map(vector_magnitude), @@ -224,6 +228,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// ``` fn visit_prepared_into( self, + params: Sink::Params, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, ) -> VortexResult