From 7dabd803ec8c348de449345f873154a76ac440b3 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 15:48:46 -0400 Subject: [PATCH 1/6] fix(array): allocate RowFn outputs with the execution allocator Signed-off-by: Connor Tsui --- vortex-array/src/memory.rs | 66 +++ .../src/scalar_fn/unstable/row/batch/tests.rs | 74 +++- .../src/scalar_fn/unstable/row/execute/mod.rs | 3 + .../scalar_fn/unstable/row/execute/owned.rs | 30 +- .../unstable/row/execute/packed_bool.rs | 8 +- .../scalar_fn/unstable/row/execute/retry.rs | 6 +- .../scalar_fn/unstable/row/execute/sink.rs | 13 +- .../scalar_fn/unstable/row/execute/tests.rs | 393 ++++++++++++++++++ .../unstable/row/types/element/bool.rs | 32 +- .../unstable/row/types/element/output.rs | 19 +- .../unstable/row/types/element/primitive.rs | 6 +- .../row/types/sink/fixed_size_list.rs | 19 +- .../scalar_fn/unstable/row/types/sink/mod.rs | 11 +- .../unstable/row/types/sink/uninit_element.rs | 18 +- .../scalar_fn/unstable/row/types/sink/utf8.rs | 30 +- vortex-spatial/src/extension/polygon.rs | 26 +- vortex-spatial/src/scalar_fn/row.rs | 11 +- 17 files changed, 686 insertions(+), 79 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/tests.rs diff --git a/vortex-array/src/memory.rs b/vortex-array/src/memory.rs index aee3d2600b4..7cac795da2c 100644 --- a/vortex-array/src/memory.rs +++ b/vortex-array/src/memory.rs @@ -79,6 +79,7 @@ pub(crate) mod test_allocator { use std::alloc::Layout; use std::ptr::NonNull; use std::sync::Arc; + use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -103,6 +104,71 @@ pub(crate) mod test_allocator { } } + /// Tracks live allocation ranges so tests can identify the returned payload's allocator. + #[derive(Clone, Debug, Default)] + pub(crate) struct AllocationTracker { + allocations: Arc>>, + } + + impl AllocationTracker { + #[track_caller] + pub(crate) fn assert_owns(&self, values: &[T]) { + assert!( + !values.is_empty(), + "allocation ownership needs a nonempty payload", + ); + let start = values.as_ptr() as usize; + let end = start + std::mem::size_of_val(values); + assert!( + self.allocations + .lock() + .unwrap() + .iter() + .any(|&(base, size)| base <= start && end <= base + size), + "the returned payload must be backed by the configured allocator", + ); + } + + pub(crate) fn live_allocations(&self) -> usize { + self.allocations.lock().unwrap().len() + } + } + + #[derive(Debug)] + struct TrackingAllocator(AllocationTracker); + + // SAFETY: allocation and deallocation use Global with the original layout. The tracker only + // records address ranges and never accesses their contents. + unsafe impl Allocator for TrackingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + let allocation = Global.allocate(layout)?; + self.0 + .allocations + .lock() + .unwrap() + .push((allocation.cast::().as_ptr() as usize, allocation.len())); + Ok(allocation) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + self.0 + .allocations + .lock() + .unwrap() + .retain(|&(base, _)| base != ptr.as_ptr() as usize); + // SAFETY: this allocation came from Global with the supplied layout. + unsafe { Global.deallocate(ptr, layout) } + } + } + + pub(crate) fn tracking_allocator() -> (BufferAllocatorRef, AllocationTracker) { + let tracker = AllocationTracker::default(); + ( + BufferAllocatorRef::new(TrackingAllocator(tracker.clone())), + tracker, + ) + } + pub(crate) fn counting_allocator() -> (BufferAllocatorRef, Arc) { let allocations = Arc::new(AtomicUsize::new(0)); ( 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 9a1c91dfd38..a3c5dc1d446 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -7,6 +7,7 @@ use std::sync::atomic::Ordering; use rstest::rstest; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -21,6 +22,7 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; +use crate::arrays::Constant; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; @@ -35,6 +37,7 @@ use crate::dtype::Nullability; use crate::dtype::extension::ExtDTypeRef; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; +use crate::memory::test_allocator::tracking_allocator; use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; @@ -47,6 +50,7 @@ use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::RowVisitor; use crate::scalar_fn::unstable::row::Utf8Column; +use crate::scalar_fn::unstable::row::Utf8Sink; use crate::scalar_fn::unstable::row::execute::execute_bool_dense_attempt; use crate::scalar_fn::unstable::row::execute_rows; use crate::scalar_fn::unstable::row::row_fn_return_dtype; @@ -253,11 +257,12 @@ impl OutputElement for NullProducingI64 { DType::from(i64::PTYPE) } - fn build(values: Vec) -> ArrayRef { - let values: Vec<_> = values.into_iter().map(|value| value.0).collect(); + fn build(values: BufferMut, allocator: &BufferAllocatorRef) -> ArrayRef { + let mut output = allocator.with_capacity(values.len()); + output.extend(values.iter().map(|value| value.0)); let validity = Validity::from_iter((0..values.len()).map(|index| index != 0)); - PrimitiveArray::new(values, validity).into_array() + PrimitiveArray::new(output.freeze(), validity).into_array() } } @@ -276,8 +281,12 @@ unsafe impl OutputSink for I64Sink { DType::from(i64::PTYPE) } - fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { - Ok(Self(BufferMut::zeroed(rows))) + fn with_capacity( + rows: usize, + _params: &Self::Params, + allocator: &BufferAllocatorRef, + ) -> VortexResult { + Ok(Self(allocator.zeroed(rows))) } fn rows(&mut self) -> Self::Rows<'_> { @@ -1023,12 +1032,18 @@ fn test_deferred_owned_execution_retries_null_row_failure() -> VortexResult<()> let lhs = PrimitiveArray::new(vec![1_i64, i64::MAX], validity.clone()).into_array(); let rhs = ConstantArray::new(1_i64, 2).into_array(); let args = VecExecutionArgs::new(vec![lhs, rhs], 2); - let mut ctx = array_session().create_execution_ctx(); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .create_execution_ctx() + .with_allocator(allocator); let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; let expected = PrimitiveArray::new(vec![2_i64, 0], validity).into_array(); - assert_arrays_eq!(&actual, &expected, &mut ctx); + // Canonical masking retains the primitive payload without allocating replacement values. + let actual = actual.execute::(&mut ctx)?; + tracker.assert_owns(actual.as_slice::()); + assert_arrays_eq!(actual.as_ref(), &expected, &mut ctx); assert_eq!(function.prepare_count(), 2); Ok(()) } @@ -1450,3 +1465,48 @@ fn test_undeclared_output_dtype_keeps_the_storage_dtype() -> VortexResult<()> { assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } + +#[derive(Clone)] +struct ConstantString; + +impl RowFn for ConstantString { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.constant_string"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64,), Utf8Sink, _>((), |_, output| { + output.write("an external UTF-8 payload"); + }) + } +} + +#[test] +fn constant_output_retains_execution_allocator_payload() -> VortexResult<()> { + let input = ConstantArray::new(1_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .create_execution_ctx() + .with_allocator(allocator); + + let output = execute_rows(&ConstantString, &EmptyOptions, &args, &mut ctx)?; + let constant = output.as_::(); + let scalar = constant.scalar(); + let value = scalar.as_utf8().value().unwrap(); + assert_eq!(value.as_str(), "an external UTF-8 payload"); + assert_eq!(output.len(), 3); + tracker.assert_owns(value.inner().as_slice()); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs index 3c1fa4fef50..c9e17da1d09 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -27,3 +27,6 @@ mod sink; pub(super) use sink::execute_sink; pub(super) use sink::execute_sink_filtered; pub(super) use sink::execute_sink_valid_rows; + +#[cfg(test)] +mod tests; 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 bab221c9b1f..2e28143534f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -3,7 +3,7 @@ //! Executes row kernels that return one independent owned value per row. //! -//! [`execute_owned`] writes fallible row results into spare vector capacity and reduces compact +//! [`execute_owned`] writes fallible row results into spare buffer capacity and reduces compact //! failure evidence outside the hot loop. [`execute_owned_infallible`] lets the output type map a //! validated row source directly into its physical representation. The `_valid_rows` variants skip //! invalid rows over the original inputs, and the `_filtered` variants read inputs filtered to the @@ -56,9 +56,11 @@ where vortex_bail!("a decoded row input does not address exactly {row_count} rows"); }; - Ok(Out::build_from(source, |elements| { - apply(&prepared, elements) - })) + Ok(Out::build_from( + source, + |elements| apply(&prepared, elements), + ctx.allocator(), + )) } /// Decode nullable inputs, then store one output for each valid row from an infallible kernel. @@ -139,9 +141,8 @@ where let prepared = prepare(Args::const_values(&columns)); let valid_rows = valid.bit_buffer(); - let mut values: Vec = std::iter::repeat_with(Out::default) - .take(valid_rows.len()) - .collect(); + let mut values = ctx.allocator().with_capacity::(valid_rows.len()); + values.extend(std::iter::repeat_with(Out::default).take(valid_rows.len())); let mut failure = Fail::default(); let mut filtered_index = 0; @@ -180,7 +181,7 @@ where finish_failure(failure)?; - Ok(Out::build(values)) + Ok(Out::build(values, ctx.allocator())) } /// Decode nullable inputs, then store outputs and combine failure evidence for valid rows. @@ -213,9 +214,8 @@ where ); let prepared = prepare(Args::const_values(&columns)); - let mut values: Vec = std::iter::repeat_with(Out::default) - .take(row_count) - .collect(); + let mut values = ctx.allocator().with_capacity::(row_count); + values.extend(std::iter::repeat_with(Out::default).take(row_count)); let mut failure = Fail::default(); if let Some(views) = Args::views_if_no_consts(&columns) { @@ -251,7 +251,7 @@ where finish_failure(failure)?; - Ok(Some(Out::build(values))) + Ok(Some(Out::build(values, ctx.allocator()))) } /// Decode every input column, then store outputs and combine per-row failure evidence. @@ -267,7 +267,7 @@ where Out: OutputElement, Fail: FailureEvidence, { - // The output vector stays at length zero until every slot is initialized so that an unwind + // The output buffer 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::() }; @@ -276,7 +276,7 @@ where let prepared = prepare(Args::const_values(&columns)); let row_count = args.row_count(); - let mut values = Vec::::with_capacity(row_count); + let mut values = ctx.allocator().with_capacity::(row_count); let output = &mut values.spare_capacity_mut()[..row_count]; let Some(source) = decoded_source::(&columns, row_count) else { @@ -291,5 +291,5 @@ where // Defer rich error construction until after the row loop. finish_failure(failure)?; - Ok(Out::build(values)) + Ok(Out::build(values, ctx.allocator())) } diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs index f60958dabe5..d494558d0ff 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs @@ -50,9 +50,9 @@ where apply(unsafe { source.get_unchecked(index) }) }; let values = if MULTIVERSIONED { - BitBuffer::collect_bool_multiversioned(row_count, collect) + BitBuffer::collect_bool_multiversioned_in(row_count, collect, ctx.allocator().clone()) } else { - BitBuffer::collect_bool(row_count, collect) + BitBuffer::collect_bool_in(row_count, collect, ctx.allocator().clone()) }; Ok(BoolArray::new(values, Validity::NonNullable).into_array()) @@ -97,9 +97,9 @@ where }; let values = if MULTIVERSIONED { - BitBuffer::collect_bool_multiversioned(row_count, collect) + BitBuffer::collect_bool_multiversioned_in(row_count, collect, ctx.allocator().clone()) } else { - BitBuffer::collect_bool(row_count, collect) + BitBuffer::collect_bool_in(row_count, collect, ctx.allocator().clone()) }; finish_failure(failure)?; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs index 13cc8734f56..a843ef22e6c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs @@ -51,7 +51,7 @@ where Out: OutputElement, Fail: FailureEvidence, { - // The output vector stays at length zero until every slot is initialized so that an unwind + // The output buffer 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::() }; @@ -62,7 +62,7 @@ where let prepared = prepare(Args::const_values(&columns)); let row_count = args.row_count(); - let mut values = Vec::::with_capacity(row_count); + let mut values = ctx.allocator().with_capacity::(row_count); let output = &mut values.spare_capacity_mut()[..row_count]; let failure_evidence = if let Some(views) = Args::views_if_no_consts(&columns) { @@ -106,7 +106,7 @@ where unsafe { values.set_len(row_count) }; match finish_failure(failure_evidence) { - Ok(()) => Ok(DenseAttempt::Values(Out::build(values))), + Ok(()) => Ok(DenseAttempt::Values(Out::build(values, ctx.allocator()))), Err(error) => Ok(DenseAttempt::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 8d9ae0f103a..000521ff8d6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -45,7 +45,7 @@ where let const_values = Args::const_values(&columns); let prepared = prepare(const_values); - let mut sink = Sink::with_capacity(row_count, params)?; + let mut sink = Sink::with_capacity(row_count, params, ctx.allocator())?; // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. { @@ -209,7 +209,7 @@ where ); let original_len = valid.len(); - let mut sink = Sink::with_capacity(original_len, params)?; + let mut sink = Sink::with_capacity(original_len, params, ctx.allocator())?; let valid_rows = valid.bit_buffer(); let views = Args::views_if_no_consts(&columns); @@ -318,7 +318,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, params)?; + let sink = Sink::with_capacity(row_count, params, ctx.allocator())?; let valid_rows = valid.bit_buffer(); vortex_ensure_eq!( @@ -338,6 +338,7 @@ where #[cfg(test)] mod tests { + use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; @@ -372,7 +373,11 @@ mod tests { DType::from(i64::PTYPE) } - fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { + fn with_capacity( + rows: usize, + _params: &Self::Params, + _allocator: &BufferAllocatorRef, + ) -> VortexResult { Ok(Self(vec![0; rows])) } diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs new file mode 100644 index 00000000000..d659253cf19 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -0,0 +1,393 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Verifies ownership of RowFn output allocations independently of input decoding. + +use rstest::rstest; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_mask::MaskValuesRef; + +use super::DenseAttempt; +use super::execute_owned; +use super::execute_owned_bool; +use super::execute_owned_dense_attempt; +use super::execute_owned_infallible; +use super::execute_owned_infallible_bool; +use super::execute_owned_infallible_filtered; +use super::execute_owned_infallible_valid_rows; +use super::execute_sink; +use super::execute_sink_filtered; +use super::execute_sink_valid_rows; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::Bool; +use crate::arrays::ConstantArray; +use crate::arrays::FixedSizeList; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinView; +use crate::arrays::bool::BoolArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::memory::MemorySessionExt; +use crate::memory::test_allocator::tracking_allocator; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::FixedSizeListSink; +use crate::scalar_fn::unstable::row::InitializedElement; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::UninitElementSink; +use crate::scalar_fn::unstable::row::Utf8Sink; + +#[derive(Clone, Copy)] +enum Traversal { + Infallible, + Fallible, + DenseAttempt, + Selected, + Filtered, +} + +fn collect_owned( + traversal: Traversal, + args: &VecExecutionArgs, + valid: &MaskValuesRef, + ctx: &mut ExecutionCtx, + apply: impl Fn(i64) -> Out, +) -> VortexResult { + match traversal { + Traversal::Infallible => execute_owned_infallible::<(i64,), Out, ()>( + args, + ctx, + |_| (), + |_, (value,)| apply(value), + ), + Traversal::Fallible => execute_owned::<(i64,), Out, (), bool>( + args, + ctx, + |_| (), + |_, (value,)| (apply(value), false), + |_| Ok(()), + ), + Traversal::DenseAttempt => { + match execute_owned_dense_attempt::<(i64,), Out, (), bool>( + args, + ctx, + |_| (), + |_, (value,)| (apply(value), false), + |_| Ok(()), + )? { + DenseAttempt::Values(values) => Ok(values), + DenseAttempt::DeferredError(error) => Err(error), + } + } + Traversal::Selected => execute_owned_infallible_valid_rows::<(i64,), Out, ()>( + args, + valid, + ctx, + |_| (), + |_, (value,)| apply(value), + )? + .ok_or_else(|| vortex_err!("canonical input must support selected rows")), + Traversal::Filtered => execute_owned_infallible_filtered::<(i64,), Out, ()>( + args, + valid, + ctx, + |_| (), + |_, (value,)| apply(value), + ), + } +} + +fn canonical_args(traversal: Traversal, constant: bool) -> VecExecutionArgs { + let values = if matches!(traversal, Traversal::Filtered) { + vec![2_i64, 4] + } else { + vec![2_i64, 3, 4] + }; + let len = values.len(); + let input = if constant { + ConstantArray::new(2_i64, len).into_array() + } else { + PrimitiveArray::from_iter(values).into_array() + }; + VecExecutionArgs::new(vec![input], len) +} + +fn selected_rows() -> MaskValuesRef { + let Mask::Values(valid) = Mask::from_iter([true, false, true]) else { + unreachable!("the test mask is partially valid"); + }; + valid +} + +#[rstest] +#[case::infallible(Traversal::Infallible)] +#[case::fallible(Traversal::Fallible)] +#[case::dense_attempt(Traversal::DenseAttempt)] +#[case::selected(Traversal::Selected)] +#[case::filtered(Traversal::Filtered)] +fn owned_payload_uses_context_allocator( + #[case] traversal: Traversal, + #[values(false, true)] constant: bool, + #[values(false, true)] boolean: bool, +) -> VortexResult<()> { + // Prepare canonical inputs and the selection before creating either measured allocator. + let args = canonical_args(traversal, constant); + let valid = selected_rows(); + let (session_allocator, session_tracker) = tracking_allocator(); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .with_allocator(session_allocator) + .create_execution_ctx() + .with_allocator(allocator); + + let output = if boolean { + collect_owned(traversal, &args, &valid, &mut ctx, |value| value % 2 == 0)? + } else { + collect_owned(traversal, &args, &valid, &mut ctx, |value| value)? + }; + if boolean { + tracker.assert_owns(output.as_::().to_bit_buffer().inner().as_slice()); + } else { + tracker.assert_owns(output.as_::().as_slice::()); + } + assert_eq!(session_tracker.live_allocations(), 0); + assert_eq!(tracker.live_allocations(), 1); + drop(output); + assert_eq!(tracker.live_allocations(), 0); + Ok(()) +} + +#[rstest] +fn packed_boolean_payload_uses_allocator( + #[values(false, true)] multiversioned: bool, + #[values(false, true)] deferred: bool, +) -> VortexResult<()> { + let args = canonical_args(Traversal::Infallible, false); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .create_execution_ctx() + .with_allocator(allocator); + + let output = match (multiversioned, deferred) { + (false, false) => { + execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 2)? + } + (true, false) => { + execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, |(value,)| value > 2)? + } + (false, true) => execute_owned_bool::<(i64,), (), bool, false>( + &args, + &mut ctx, + |_| (), + |_, (value,)| (value > 2, false), + |_| Ok(()), + )?, + (true, true) => execute_owned_bool::<(i64,), (), bool, true>( + &args, + &mut ctx, + |_| (), + |_, (value,)| (value > 2, false), + |_| Ok(()), + )?, + }; + tracker.assert_owns(output.as_::().to_bit_buffer().inner().as_slice()); + assert_eq!(tracker.live_allocations(), 1); + Ok(()) +} + +fn collect_sink( + traversal: Traversal, + args: &VecExecutionArgs, + valid: &MaskValuesRef, + params: &Sink::Params, + ctx: &mut ExecutionCtx, + apply: impl Fn(Sink::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Sink: OutputSink, + ApplyResult: SinkResult, +{ + match traversal { + Traversal::Infallible => execute_sink::<(i64,), (), Sink, ApplyResult>( + args, + params, + ctx, + |_| (), + |_, _, row| apply(row), + ), + Traversal::Selected => execute_sink_valid_rows::<(i64,), (), Sink, ApplyResult>( + args, + valid, + params, + ctx, + |_| (), + |_, _, row| apply(row), + )? + .ok_or_else(|| vortex_err!("canonical input must support selected rows")), + Traversal::Filtered => execute_sink_filtered::<(i64,), (), Sink, ApplyResult>( + args, + valid, + params, + ctx, + |_| (), + |_, _, row| apply(row), + ), + _ => vortex_bail!("this test traversal requires an owned output"), + } +} + +#[rstest] +#[case::dense(Traversal::Infallible)] +#[case::selected(Traversal::Selected)] +#[case::filtered(Traversal::Filtered)] +fn sink_payloads_use_allocator(#[case] traversal: Traversal) -> VortexResult<()> { + let args = canonical_args(traversal, false); + let valid = selected_rows(); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .create_execution_ctx() + .with_allocator(allocator); + + let scalar = collect_sink::, _>( + traversal, + &args, + &valid, + &(), + &mut ctx, + |row| { + // SAFETY: writes the supplied slot and returns its token without modifying the slot. + unsafe { InitializedElement::write(row, 42) } + }, + )?; + tracker.assert_owns(scalar.as_::().as_slice::()); + + let strings = collect_sink::(traversal, &args, &valid, &(), &mut ctx, |row| { + row.write("an external UTF-8 payload") + })?; + let strings = strings.as_::(); + tracker.assert_owns(strings.views()); + assert!(!strings.data_buffers().is_empty()); + for buffer in strings.data_buffers().iter() { + tracker.assert_owns(buffer.as_host().as_slice()); + } + Ok(()) +} + +#[test] +fn primitive_build_reuses_allocation() { + let (allocator, tracker) = tracking_allocator(); + let values = allocator.copy_from([1_i64, 2, 3]); + let ptr = values.as_ptr(); + let output = i64::build(values, &allocator); + assert_eq!(output.as_::().as_slice::().as_ptr(), ptr); + tracker.assert_owns(output.as_::().as_slice::()); +} + +#[test] +fn empty_outputs_and_zero_width_rows_do_not_allocate_payloads() -> VortexResult<()> { + let empty = PrimitiveArray::from_iter(std::iter::empty::()).into_array(); + let args = VecExecutionArgs::new(vec![empty], 0); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .create_execution_ctx() + .with_allocator(allocator); + + let primitive = + execute_owned_infallible::<(i64,), i64, ()>(&args, &mut ctx, |_| (), |_, (value,)| value)?; + let boolean = + execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 0)?; + let scalar = execute_sink::<(i64,), (), UninitElementSink, _>( + &args, + &(), + &mut ctx, + |_| (), + |_, _, row| { + // SAFETY: writes the supplied slot and immediately returns its token. + unsafe { InitializedElement::write(row, 0) } + }, + )?; + let mut lists = FixedSizeListSink::::with_capacity(3, &0, ctx.allocator())?; + FixedSizeListSink::::initialize_skipped_rows(&mut lists.rows()); + // SAFETY: the skipped-row initializer completed for every zero-width row. + let lists = unsafe { lists.finish() }?; + let strings = execute_sink::<(i64,), (), Utf8Sink, _>( + &args, + &(), + &mut ctx, + |_| (), + |_, _, row| row.write("unused"), + )?; + + assert!(primitive.is_empty()); + assert!(boolean.is_empty()); + assert!(scalar.is_empty()); + assert!(strings.is_empty()); + assert_eq!(lists.len(), 3); + assert!(lists.as_::().elements().is_empty()); + assert_eq!(tracker.live_allocations(), 0); + Ok(()) +} + +#[test] +fn boolean_sinks_allocate_packed_payloads_with_context_allocator() -> VortexResult<()> { + let args = canonical_args(Traversal::Infallible, false); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .create_execution_ctx() + .with_allocator(allocator); + + let scalar = execute_sink::<(i64,), (), UninitElementSink, _>( + &args, + &(), + &mut ctx, + |_| (), + |_, _, row| { + // SAFETY: writes the supplied slot and immediately returns its token. + unsafe { InitializedElement::write(row, true) } + }, + )?; + tracker.assert_owns(scalar.as_::().to_bit_buffer().inner().as_slice()); + + let mut lists = FixedSizeListSink::::with_capacity(3, &2, ctx.allocator())?; + FixedSizeListSink::::initialize_skipped_rows(&mut lists.rows()); + // SAFETY: the skipped-row initializer initialized every element. + let lists = unsafe { lists.finish() }?; + tracker.assert_owns( + lists + .as_::() + .elements() + .as_::() + .to_bit_buffer() + .inner() + .as_slice(), + ); + assert_eq!(tracker.live_allocations(), 2); + Ok(()) +} + +#[test] +fn fixed_size_list_payload_uses_allocator() -> VortexResult<()> { + let (allocator, tracker) = tracking_allocator(); + let mut sink = FixedSizeListSink::::with_capacity(3, &2, &allocator)?; + FixedSizeListSink::::initialize_skipped_rows(&mut sink.rows()); + // SAFETY: the skipped-row initializer initialized every element. + let output = unsafe { sink.finish() }?; + + tracker.assert_owns( + output + .as_::() + .elements() + .as_::() + .as_slice::(), + ); + assert_eq!(tracker.live_allocations(), 1); + Ok(()) +} 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 a1c1fc61934..28bb5c54d1f 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 @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSource; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -94,22 +96,34 @@ impl OutputElement for bool { DType::Bool(Nullability::NonNullable) } - fn build(values: Vec) -> ArrayRef { - // `From>` uses the bulk bit-packing path. - BoolArray::new(BitBuffer::from(values), Validity::NonNullable).into_array() + fn build(values: BufferMut, allocator: &BufferAllocatorRef) -> ArrayRef { + let values = values.as_slice(); + let packed = BitBuffer::collect_bool_multiversioned_in( + values.len(), + |index| { + // SAFETY: the collector only requests indices below `values.len()`. + unsafe { *values.get_unchecked(index) } + }, + allocator.clone(), + ); + BoolArray::new(packed, Validity::NonNullable).into_array() } - fn build_from(source: S, apply: F) -> ArrayRef + fn build_from(source: S, apply: F, allocator: &BufferAllocatorRef) -> ArrayRef where S: IndexedSource, F: Fn(S::Item) -> Self, { let len = source.len(); - let values = BitBuffer::collect_bool(len, |index| { - // SAFETY: `collect_bool` only invokes this closure with `index < len`, and - // `len` is `source.len()`. - apply(unsafe { source.get_unchecked(index) }) - }); + let values = BitBuffer::collect_bool_in( + len, + |index| { + // SAFETY: `collect_bool_in` only invokes this closure with `index < len`, and + // `len` is `source.len()`. + apply(unsafe { source.get_unchecked(index) }) + }, + allocator.clone(), + ); BoolArray::new(values, Validity::NonNullable).into_array() } 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 3b92107931b..28f3d9db83d 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 @@ -5,6 +5,8 @@ //! //! [`OutputElement`] describes fixed-dtype values returned independently by each row invocation. +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSource; use vortex_compute::lane_kernels::IndexedSourceExt; @@ -14,7 +16,8 @@ use crate::dtype::DType; /// An owned row value that can be built into an all-valid column. /// /// Skip-invalid execution uses [`Default`] only as a placeholder for invalid rows. Batch execution -/// masks those rows before returning the output. +/// masks those rows before returning the output. Element types must have nonzero size because +/// collection uses [`BufferMut`]. pub trait OutputElement: 'static + Sized + Default { /// The dtype of columns built from this element type. **Must** be non-nullable: nullability is /// derived from the inputs by batch execution. @@ -30,27 +33,31 @@ pub trait OutputElement: 'static + Sized + Default { /// The returned column must contain `values.len()` rows and match /// [`element_dtype`](Self::element_dtype) except for outer nullability. The default /// [`build_from`](Self::build_from) implementation and valid-row execution call this method. - fn build(values: Vec) -> ArrayRef; + /// + /// Reuse `values` when it already has the required physical representation. Any new payload + /// buffers must use `allocator`. + fn build(values: BufferMut, allocator: &BufferAllocatorRef) -> ArrayRef; /// Map a contiguous row source directly into an all-valid column. /// - /// The default collects one value per row into a [`Vec`] before calling [`build`](Self::build). + /// The default collects into a [`BufferMut`] using `allocator`, then calls [`build`](Self::build). /// An output type can override this method when its physical representation supports a more /// efficient bulk mapping. The implementation **must** call `apply` exactly once for every /// source row in increasing order and return the same values as the default implementation. + /// Any new payload buffers must use `allocator`. /// /// An override **must not** introduce value-dependent errors or panics. Fallible operations /// must use a fallible visitor path so that [`RowFn::INFALLIBLE`] continues to protect optimizer /// transformations. /// /// [`RowFn::INFALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::INFALLIBLE - fn build_from(source: S, apply: F) -> ArrayRef + fn build_from(source: S, apply: F, allocator: &BufferAllocatorRef) -> ArrayRef where S: IndexedSource, F: Fn(S::Item) -> Self, { let row_count = source.len(); - let mut values = Vec::::with_capacity(row_count); + let mut values = allocator.with_capacity::(row_count); let output = &mut values.spare_capacity_mut()[..row_count]; source.map_into(output, apply); @@ -58,6 +65,6 @@ pub trait OutputElement: 'static + Sized + Default { // SAFETY: normal completion of `map_into` initializes every output slot exactly once. unsafe { values.set_len(row_count) }; - Self::build(values) + Self::build(values, allocator) } } 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 32885112b51..971b1add143 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 @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; @@ -102,7 +104,7 @@ impl OutputElement for T { DType::Primitive(T::PTYPE, Nullability::NonNullable) } - fn build(values: Vec) -> ArrayRef { - PrimitiveArray::new(values, Validity::NonNullable).into_array() + fn build(values: BufferMut, _allocator: &BufferAllocatorRef) -> ArrayRef { + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array() } } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs index 856602db0a4..e266e0e986d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs @@ -9,6 +9,8 @@ use std::mem::MaybeUninit; use std::sync::Arc; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -87,11 +89,11 @@ impl FillDefault for FixedSizeRows<'_, T> { /// [`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, - + values: BufferMut, + /// Allocator for any physical conversion when the sink finishes. + allocator: BufferAllocatorRef, /// The number of elements in each output row. width: usize, - /// The number of output rows. row_count: usize, } @@ -115,7 +117,11 @@ unsafe impl OutputSink for FixedSizeListSink< ) } - fn with_capacity(rows: usize, params: &Self::Params) -> VortexResult { + fn with_capacity( + rows: usize, + params: &Self::Params, + allocator: &BufferAllocatorRef, + ) -> VortexResult { let width = *params; let element_capacity = rows.checked_mul(width).ok_or_else(|| { vortex_err!( @@ -125,7 +131,8 @@ unsafe impl OutputSink for FixedSizeListSink< })?; Ok(Self { - values: Vec::with_capacity(element_capacity), + values: allocator.with_capacity(element_capacity), + allocator: allocator.clone(), width, row_count: rows, }) @@ -155,7 +162,7 @@ unsafe impl OutputSink for FixedSizeListSink< // `row_count * width` elements. unsafe { self.values.set_len(element_count) }; - let elements = T::build(self.values); + let elements = T::build(self.values, &self.allocator); let lists = FixedSizeListArray::new( elements, fixed_size_list_size(self.width), diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs index 5eb045f958d..e35ea66a481 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/mod.rs @@ -7,6 +7,7 @@ //! uninitialized scalar storage, [`FixedSizeListSink`] provides runtime-width row storage, and //! [`Utf8Sink`] provides variable-length UTF-8 storage. +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use crate::ArrayRef; @@ -129,8 +130,14 @@ pub unsafe trait OutputSink: 'static + Sized { /// result, and masks the null rows. fn storage_dtype(params: &Self::Params) -> DType; - /// Allocate a sink for `rows` rows. - fn with_capacity(rows: usize, params: &Self::Params) -> VortexResult; + /// Allocate a sink for `rows` rows, using `allocator` for new output payload buffers. + /// + /// The allocator is an execution resource, separate from physical storage parameters. + fn with_capacity( + rows: usize, + params: &Self::Params, + allocator: &BufferAllocatorRef, + ) -> VortexResult; /// Borrow all output rows for the hot loop. fn rows(&mut self) -> Self::Rows<'_>; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs index 4b29d70b104..0f5bbac9e73 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs @@ -8,6 +8,8 @@ use std::mem::MaybeUninit; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use super::OutputSink; @@ -60,8 +62,9 @@ impl InitializedElement { /// initialized spare-capacity elements require no destruction. pub struct UninitElementSink { /// Spare storage written in increasing row order. - values: Vec, - + values: BufferMut, + /// Allocator for any physical conversion when the sink finishes. + allocator: BufferAllocatorRef, /// The number of slots exposed to the row loop and initialized before finishing. row_count: usize, } @@ -102,9 +105,14 @@ unsafe impl OutputSink for UninitElementSink< T::element_dtype() } - fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { + fn with_capacity( + rows: usize, + _params: &Self::Params, + allocator: &BufferAllocatorRef, + ) -> VortexResult { Ok(Self { - values: Vec::with_capacity(rows), + values: allocator.with_capacity(rows), + allocator: allocator.clone(), row_count: rows, }) } @@ -123,6 +131,6 @@ unsafe impl OutputSink for UninitElementSink< // `with_capacity` reserved every slot in that range. unsafe { self.values.set_len(self.row_count) }; - Ok(T::build(self.values)) + Ok(T::build(self.values, &self.allocator)) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/utf8.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/utf8.rs index e44640e8eb3..2ca2e5e13c3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/utf8.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/utf8.rs @@ -9,6 +9,7 @@ use std::sync::Arc; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; @@ -31,12 +32,14 @@ use crate::validity::Validity; pub struct Utf8Sink { views: BufferMut, buffers: Vec, + allocator: BufferAllocatorRef, } /// A borrowed view of all UTF-8 output rows. pub struct Utf8Rows<'a> { views: &'a mut [BinaryView], buffers: &'a mut Vec, + allocator: &'a BufferAllocatorRef, } impl ViewLen for Utf8Rows<'_> { @@ -54,6 +57,7 @@ impl FillDefault for Utf8Rows<'_> { pub struct Utf8Writer<'a> { view: &'a mut BinaryView, buffers: &'a mut Vec, + allocator: &'a BufferAllocatorRef, } impl Utf8Writer<'_> { @@ -70,7 +74,7 @@ impl Utf8Writer<'_> { .last() .is_none_or(|buffer| buffer.len().saturating_add(bytes.len()) > MAX_BUFFER_LEN); if needs_buffer { - self.buffers.push(ByteBufferMut::with_capacity(bytes.len())); + self.buffers.push(self.allocator.with_capacity(bytes.len())); } let buffer_index = u32::try_from(self.buffers.len() - 1) @@ -102,13 +106,18 @@ unsafe impl OutputSink for Utf8Sink { DType::Utf8(Nullability::NonNullable) } - fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult { - let mut views = BufferMut::with_capacity(rows); + fn with_capacity( + rows: usize, + _params: &Self::Params, + allocator: &BufferAllocatorRef, + ) -> VortexResult { + let mut views = allocator.with_capacity(rows); views.push_n(BinaryView::empty_view(), rows); Ok(Self { views, buffers: Vec::new(), + allocator: allocator.clone(), }) } @@ -116,6 +125,7 @@ unsafe impl OutputSink for Utf8Sink { Utf8Rows { views: self.views.as_mut_slice(), buffers: &mut self.buffers, + allocator: &self.allocator, } } @@ -126,6 +136,7 @@ unsafe impl OutputSink for Utf8Sink { Utf8Writer { view, buffers: rows.buffers, + allocator: rows.allocator, } } @@ -151,6 +162,7 @@ unsafe impl OutputSink for Utf8Sink { mod tests { use std::borrow::Cow; + use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -163,7 +175,11 @@ mod tests { fn sink_writes_owned_borrowed_and_cow_values() -> VortexResult<()> { let expected = ["short", "a referenced string", "owned", "borrowed cow"]; let referenced = String::from("a referenced string"); - let mut sink = ::with_capacity(expected.len(), &())?; + let mut sink = ::with_capacity( + expected.len(), + &(), + BufferAllocatorRef::static_ref(), + )?; { let mut rows = ::rows(&mut sink); @@ -195,12 +211,14 @@ mod tests { #[test] fn sink_finishes_empty_and_skipped_rows() -> VortexResult<()> { - let empty = ::with_capacity(0, &())?; + let empty = + ::with_capacity(0, &(), BufferAllocatorRef::static_ref())?; // SAFETY: a zero-row sink has no rows to initialize. let empty = unsafe { ::finish(empty) }?; assert!(empty.is_empty()); - let mut skipped = ::with_capacity(2, &())?; + let mut skipped = + ::with_capacity(2, &(), BufferAllocatorRef::static_ref())?; ::initialize_skipped_rows(&mut ::rows( &mut skipped, )); diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index 87b542f73f7..0a0763a7140 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -33,6 +33,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; +use vortex_array::memory::BufferAllocatorRef; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; @@ -120,11 +121,14 @@ fn polygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> Pol /// Build canonical non-nullable 2-D polygon storage from row-oriented `geo_types` polygons. pub(crate) fn build_polygon_storage( polygons: &[geo_types::Polygon], + allocator: &BufferAllocatorRef, ) -> VortexResult { - let mut xs = Vec::new(); - let mut ys = Vec::new(); - let mut ring_offsets = vec![0_u64]; - let mut polygon_offsets = vec![0_u64]; + let mut xs = allocator.with_capacity::(0); + let mut ys = allocator.with_capacity::(0); + let mut ring_offsets = allocator.with_capacity::(1); + ring_offsets.push(0); + let mut polygon_offsets = allocator.with_capacity::(polygons.len() + 1); + polygon_offsets.push(0); for polygon in polygons { let exterior = (!polygon.exterior().is_empty()).then_some(polygon.exterior()); @@ -143,19 +147,25 @@ pub(crate) fn build_polygon_storage( } let coordinates = StructArray::from_fields(&[ - ("x", PrimitiveArray::from_iter(xs).into_array()), - ("y", PrimitiveArray::from_iter(ys).into_array()), + ( + "x", + PrimitiveArray::new(xs.freeze(), Validity::NonNullable).into_array(), + ), + ( + "y", + PrimitiveArray::new(ys.freeze(), Validity::NonNullable).into_array(), + ), ])? .into_array(); let rings = ListArray::try_new( coordinates, - PrimitiveArray::from_iter(ring_offsets).into_array(), + PrimitiveArray::new(ring_offsets.freeze(), Validity::NonNullable).into_array(), Validity::NonNullable, )? .into_array(); let storage = ListArray::try_new( rings, - PrimitiveArray::from_iter(polygon_offsets).into_array(), + PrimitiveArray::new(polygon_offsets.freeze(), Validity::NonNullable).into_array(), Validity::NonNullable, )? .into_array(); diff --git a/vortex-spatial/src/scalar_fn/row.rs b/vortex-spatial/src/scalar_fn/row.rs index c0255d44862..b3f1cf8d80b 100644 --- a/vortex-spatial/src/scalar_fn/row.rs +++ b/vortex-spatial/src/scalar_fn/row.rs @@ -11,6 +11,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::memory::BufferAllocatorRef; use vortex_array::scalar_fn::unstable::row::InputElement; use vortex_array::scalar_fn::unstable::row::OutputSink; use vortex_array::scalar_fn::unstable::row::Preinitialized; @@ -136,6 +137,7 @@ unsafe impl InputElement for GeometryRow { /// Row output for native 2-D polygons. pub(crate) struct PolygonSink { polygons: Vec>, + allocator: BufferAllocatorRef, } fn empty_polygon() -> GeoPolygon { @@ -156,9 +158,14 @@ unsafe impl OutputSink for PolygonSink { polygon_storage_dtype(Dimension::Xy, Nullability::NonNullable) } - fn with_capacity(rows: usize, (): &Self::Params) -> VortexResult { + fn with_capacity( + rows: usize, + (): &Self::Params, + allocator: &BufferAllocatorRef, + ) -> VortexResult { Ok(Self { polygons: vec![empty_polygon(); rows], + allocator: allocator.clone(), }) } @@ -172,6 +179,6 @@ unsafe impl OutputSink for PolygonSink { } unsafe fn finish(self) -> VortexResult { - build_polygon_storage(&self.polygons) + build_polygon_storage(&self.polygons, &self.allocator) } } From 7bfdcad108787a4ec220dabd4dd3abf938e62b5b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 16:24:01 -0400 Subject: [PATCH 2/6] refactor(array): let RowFn outputs own their collection storage Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 27 ++++++- .../scalar_fn/unstable/row/execute/owned.rs | 51 +++++++------ .../scalar_fn/unstable/row/execute/retry.rs | 19 ++--- .../scalar_fn/unstable/row/execute/tests.rs | 73 ++++++++++++++++++- .../src/scalar_fn/unstable/row/mod.rs | 1 + .../unstable/row/types/element/bool.rs | 43 ++++++++--- .../unstable/row/types/element/mod.rs | 1 + .../unstable/row/types/element/output.rs | 60 ++++++++++----- .../unstable/row/types/element/primitive.rs | 25 ++++++- .../src/scalar_fn/unstable/row/types/mod.rs | 1 + .../row/types/sink/fixed_size_list.rs | 22 +++--- .../unstable/row/types/sink/uninit_element.rs | 26 +++---- 12 files changed, 255 insertions(+), 94 deletions(-) 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 a3c5dc1d446..4baae4df3d3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -45,6 +46,7 @@ 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::InputElement; +use crate::scalar_fn::unstable::row::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; @@ -253,14 +255,31 @@ unsafe impl InputElement for DenseRetryI64 { struct NullProducingI64(i64); impl OutputElement for NullProducingI64 { + type Buffer = BufferMut; + fn element_dtype() -> DType { DType::from(i64::PTYPE) } - fn build(values: BufferMut, allocator: &BufferAllocatorRef) -> ArrayRef { - let mut output = allocator.with_capacity(values.len()); - output.extend(values.iter().map(|value| value.0)); - let validity = Validity::from_iter((0..values.len()).map(|index| index != 0)); + fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { + allocator.with_capacity(rows) + } +} + +// SAFETY: clearing the length preserves the slots, and these values require no destruction. +unsafe impl OutputBuffer for BufferMut { + fn slots(&mut self) -> &mut [MaybeUninit] { + self.clear(); + self.spare_capacity_mut() + } + + unsafe fn finish(mut self, len: usize, allocator: &BufferAllocatorRef) -> ArrayRef { + // SAFETY: the caller initialized the first `len` slots. + unsafe { self.set_len(len) }; + + let mut output = allocator.with_capacity(len); + output.extend(self.iter().map(|value| value.0)); + let validity = Validity::from_iter((0..len).map(|index| index != 0)); PrimitiveArray::new(output.freeze(), validity).into_array() } 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 2e28143534f..b34d89a391e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -23,6 +23,7 @@ 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::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::types::decoded_source; use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; @@ -141,8 +142,13 @@ where let prepared = prepare(Args::const_values(&columns)); let valid_rows = valid.bit_buffer(); - let mut values = ctx.allocator().with_capacity::(valid_rows.len()); - values.extend(std::iter::repeat_with(Out::default).take(valid_rows.len())); + let mut values = Out::allocate(valid_rows.len(), ctx.allocator()); + let output = &mut values.slots()[..valid_rows.len()]; + + for slot in output.iter_mut() { + slot.write(Out::default()); + } + let mut failure = Fail::default(); let mut filtered_index = 0; @@ -158,8 +164,8 @@ where let elements = unsafe { Args::get_from_views_unchecked(&views, filtered_index) }; let (value, row_failure) = apply(&prepared, elements); - // SAFETY: every set index is below the mask length, which sized `values`. - unsafe { *values.get_unchecked_mut(index) = value }; + // SAFETY: every set index is below the mask length, which sized `output`. + unsafe { output.get_unchecked_mut(index) }.write(value); failure |= row_failure; filtered_index += 1; }); @@ -172,8 +178,8 @@ where valid_rows.for_each_set_index(|index| { let (value, row_failure) = apply(&prepared, Args::get(&columns, filtered_index)); - // SAFETY: every set index is below the mask length, which sized `values`. - unsafe { *values.get_unchecked_mut(index) = value }; + // SAFETY: every set index is below the mask length, which sized `output`. + unsafe { output.get_unchecked_mut(index) }.write(value); failure |= row_failure; filtered_index += 1; }); @@ -181,7 +187,8 @@ where finish_failure(failure)?; - Ok(Out::build(values, ctx.allocator())) + // SAFETY: every output slot contains either its placeholder or the row result. + Ok(unsafe { values.finish(valid_rows.len(), ctx.allocator()) }) } /// Decode nullable inputs, then store outputs and combine failure evidence for valid rows. @@ -214,8 +221,13 @@ where ); let prepared = prepare(Args::const_values(&columns)); - let mut values = ctx.allocator().with_capacity::(row_count); - values.extend(std::iter::repeat_with(Out::default).take(row_count)); + let mut values = Out::allocate(row_count, ctx.allocator()); + let output = &mut values.slots()[..row_count]; + + for slot in output.iter_mut() { + slot.write(Out::default()); + } + let mut failure = Fail::default(); if let Some(views) = Args::views_if_no_consts(&columns) { @@ -231,7 +243,7 @@ where let (value, row_failure) = apply(&prepared, elements); // SAFETY: the mask length check proved that every set index is below `row_count`. - unsafe { *values.get_unchecked_mut(index) = value }; + unsafe { output.get_unchecked_mut(index) }.write(value); failure |= row_failure; }); } else { @@ -244,14 +256,15 @@ where let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); // SAFETY: the mask length check proved that every set index is below `row_count`. - unsafe { *values.get_unchecked_mut(index) = value }; + unsafe { output.get_unchecked_mut(index) }.write(value); failure |= row_failure; }); } finish_failure(failure)?; - Ok(Some(Out::build(values, ctx.allocator()))) + // SAFETY: every output slot contains either its placeholder or the row result. + Ok(Some(unsafe { values.finish(row_count, ctx.allocator()) })) } /// Decode every input column, then store outputs and combine per-row failure evidence. @@ -267,8 +280,7 @@ where Out: OutputElement, Fail: FailureEvidence, { - // The output buffer 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 + // Errors and unwinds abandon partially initialized slots. The assertion ensures that no // initialized value requires a destructor to run. const { assert_owned_output_needs_no_drop::() }; @@ -276,20 +288,17 @@ where let prepared = prepare(Args::const_values(&columns)); let row_count = args.row_count(); - let mut values = ctx.allocator().with_capacity::(row_count); - let output = &mut values.spare_capacity_mut()[..row_count]; + let mut values = Out::allocate(row_count, ctx.allocator()); + let output = &mut values.slots()[..row_count]; let Some(source) = decoded_source::(&columns, row_count) else { vortex_bail!("a decoded row input does not address exactly {row_count} rows"); }; let failure = source.map_checked_into(output, |elements| apply(&prepared, elements)); - // SAFETY: normal completion initializes `0..row_count` exactly once, and `values` was - // allocated with at least `row_count` capacity. - unsafe { values.set_len(row_count) }; - // Defer rich error construction until after the row loop. finish_failure(failure)?; - Ok(Out::build(values, ctx.allocator())) + // SAFETY: normal completion of `map_checked_into` initializes every output slot. + Ok(unsafe { values.finish(row_count, ctx.allocator()) }) } diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs index a843ef22e6c..8ca4babcff3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs @@ -18,6 +18,7 @@ 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::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; @@ -51,8 +52,7 @@ where Out: OutputElement, Fail: FailureEvidence, { - // The output buffer 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 + // Errors and unwinds abandon partially initialized slots. The assertion ensures that no // initialized value requires a destructor to run. const { assert_owned_output_needs_no_drop::() }; @@ -62,8 +62,8 @@ where let prepared = prepare(Args::const_values(&columns)); let row_count = args.row_count(); - let mut values = ctx.allocator().with_capacity::(row_count); - let output = &mut values.spare_capacity_mut()[..row_count]; + let mut values = Out::allocate(row_count, ctx.allocator()); + let output = &mut values.slots()[..row_count]; let failure_evidence = if let Some(views) = Args::views_if_no_consts(&columns) { // Keep this validation beside the views so LLVM sees their common length here. @@ -101,12 +101,13 @@ where accumulated_failure }; - // 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) }; - match finish_failure(failure_evidence) { - Ok(()) => Ok(DenseAttempt::Values(Out::build(values, ctx.allocator()))), + Ok(()) => { + // SAFETY: normal completion of either path initializes every output slot. + let output = unsafe { values.finish(row_count, ctx.allocator()) }; + + Ok(DenseAttempt::Values(output)) + } Err(error) => Ok(DenseAttempt::DeferredError(error)), } } diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index d659253cf19..96ee66bdb98 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -3,7 +3,10 @@ //! Verifies ownership of RowFn output allocations independently of input decoding. +use std::mem::MaybeUninit; + use rstest::rstest; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -34,16 +37,20 @@ use crate::arrays::PrimitiveArray; use crate::arrays::VarBinView; use crate::arrays::bool::BoolArrayExt; use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt; +use crate::assert_arrays_eq; +use crate::dtype::DType; use crate::memory::MemorySessionExt; use crate::memory::test_allocator::tracking_allocator; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::FixedSizeListSink; use crate::scalar_fn::unstable::row::InitializedElement; +use crate::scalar_fn::unstable::row::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::SinkResult; use crate::scalar_fn::unstable::row::UninitElementSink; use crate::scalar_fn::unstable::row::Utf8Sink; +use crate::validity::Validity; #[derive(Clone, Copy)] enum Traversal { @@ -282,11 +289,18 @@ fn sink_payloads_use_allocator(#[case] traversal: Traversal) -> VortexResult<()> } #[test] -fn primitive_build_reuses_allocation() { +fn primitive_finish_reuses_allocation() { let (allocator, tracker) = tracking_allocator(); - let values = allocator.copy_from([1_i64, 2, 3]); - let ptr = values.as_ptr(); - let output = i64::build(values, &allocator); + let mut values = i64::allocate(3, &allocator); + let slots = &mut values.slots()[..3]; + let ptr = slots.as_ptr().cast::(); + + for (slot, value) in slots.iter_mut().zip([1, 2, 3]) { + slot.write(value); + } + + // SAFETY: all three slots were initialized above. + let output = unsafe { values.finish(3, &allocator) }; assert_eq!(output.as_::().as_slice::().as_ptr(), ptr); tracker.assert_owns(output.as_::().as_slice::()); } @@ -391,3 +405,54 @@ fn fixed_size_list_payload_uses_allocator() -> VortexResult<()> { assert_eq!(tracker.live_allocations(), 1); Ok(()) } + +/// A zero-sized element whose collection storage does not depend on Vortex buffers. +#[derive(Clone, Copy, Default)] +struct One; + +impl OutputElement for One { + type Buffer = Vec>; + + fn element_dtype() -> DType { + i64::element_dtype() + } + + fn allocate(rows: usize, _allocator: &BufferAllocatorRef) -> Self::Buffer { + vec![MaybeUninit::uninit(); rows] + } +} + +// SAFETY: the vector retains the same slots, and MaybeUninit permits partial initialization. +unsafe impl OutputBuffer for Vec> { + fn slots(&mut self) -> &mut [MaybeUninit] { + self.as_mut_slice() + } + + unsafe fn finish(self, len: usize, allocator: &BufferAllocatorRef) -> ArrayRef { + let mut values = allocator.with_capacity(len); + values.extend(std::iter::repeat_n(1_i64, len)); + PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array() + } +} + +#[rstest] +#[case::infallible(Traversal::Infallible)] +#[case::fallible(Traversal::Fallible)] +#[case::dense_attempt(Traversal::DenseAttempt)] +#[case::selected(Traversal::Selected)] +#[case::filtered(Traversal::Filtered)] +fn zero_sized_output_uses_its_own_storage(#[case] traversal: Traversal) -> VortexResult<()> { + let args = canonical_args(traversal, false); + let valid = selected_rows(); + let expected = PrimitiveArray::from_iter([1_i64; 3]).into_array(); + let (allocator, tracker) = tracking_allocator(); + let mut ctx = array_session() + .create_execution_ctx() + .with_allocator(allocator); + + let output = collect_owned::(traversal, &args, &valid, &mut ctx, |_| One)?; + assert_arrays_eq!(&output, &expected, &mut ctx); + tracker.assert_owns(output.as_::().as_slice::()); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index 7c9896060f1..d94e399dbac 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -37,6 +37,7 @@ pub use types::IndexedElementTuple; pub use types::InitializedElement; pub use types::InitializedRow; pub use types::InputElement; +pub use types::OutputBuffer; pub use types::OutputElement; pub use types::OutputSink; pub use types::Preinitialized; 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 28bb5c54d1f..707b6e18e94 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 @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; + use vortex_buffer::BitBuffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; @@ -18,6 +20,7 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::ScalarValue; use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::validity::Validity; @@ -92,21 +95,14 @@ unsafe impl InputElement for bool { } impl OutputElement for bool { + type Buffer = BufferMut; + fn element_dtype() -> DType { DType::Bool(Nullability::NonNullable) } - fn build(values: BufferMut, allocator: &BufferAllocatorRef) -> ArrayRef { - let values = values.as_slice(); - let packed = BitBuffer::collect_bool_multiversioned_in( - values.len(), - |index| { - // SAFETY: the collector only requests indices below `values.len()`. - unsafe { *values.get_unchecked(index) } - }, - allocator.clone(), - ); - BoolArray::new(packed, Validity::NonNullable).into_array() + fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { + allocator.with_capacity(rows) } fn build_from(source: S, apply: F, allocator: &BufferAllocatorRef) -> ArrayRef @@ -128,3 +124,28 @@ impl OutputElement for bool { BoolArray::new(values, Validity::NonNullable).into_array() } } + +// SAFETY: clearing the length preserves the contents and exposes the same allocation each time. +// Booleans require no destruction when a partially initialized buffer is abandoned. +unsafe impl OutputBuffer for BufferMut { + fn slots(&mut self) -> &mut [MaybeUninit] { + self.clear(); + self.spare_capacity_mut() + } + + unsafe fn finish(mut self, len: usize, allocator: &BufferAllocatorRef) -> ArrayRef { + // SAFETY: the caller initialized the first `len` slots of this buffer's spare capacity. + unsafe { self.set_len(len) }; + + let values = self.as_slice(); + let packed = BitBuffer::collect_bool_multiversioned_in( + values.len(), + |index| { + // SAFETY: the collector only requests indices below `values.len()`. + unsafe { *values.get_unchecked(index) } + }, + allocator.clone(), + ); + BoolArray::new(packed, Validity::NonNullable).into_array() + } +} 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 aea9c9c4da7..ac2b4b408b2 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 @@ -16,6 +16,7 @@ mod input; pub use input::InputElement; mod output; +pub use output::OutputBuffer; pub use output::OutputElement; mod primitive; 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 28f3d9db83d..3840bec8b05 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 @@ -5,8 +5,9 @@ //! //! [`OutputElement`] describes fixed-dtype values returned independently by each row invocation. +use std::mem::MaybeUninit; + use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSource; use vortex_compute::lane_kernels::IndexedSourceExt; @@ -16,9 +17,12 @@ use crate::dtype::DType; /// An owned row value that can be built into an all-valid column. /// /// Skip-invalid execution uses [`Default`] only as a placeholder for invalid rows. Batch execution -/// masks those rows before returning the output. Element types must have nonzero size because -/// collection uses [`BufferMut`]. +/// masks those rows before returning the output. Each implementation owns allocation and array +/// construction. Execution only writes through its [`OutputBuffer`] slots. pub trait OutputElement: 'static + Sized + Default { + /// Storage used to collect this element type before constructing its column. + type Buffer: OutputBuffer + 'static; + /// The dtype of columns built from this element type. **Must** be non-nullable: nullability is /// derived from the inputs by batch execution. /// @@ -28,19 +32,15 @@ pub trait OutputElement: 'static + Sized + Default { /// [`OutputSink`]: crate::scalar_fn::unstable::row::OutputSink fn element_dtype() -> DType; - /// Build an all-valid column from one value per row. + /// Allocate storage with at least `rows` writable slots using the execution allocator. /// - /// The returned column must contain `values.len()` rows and match - /// [`element_dtype`](Self::element_dtype) except for outer nullability. The default - /// [`build_from`](Self::build_from) implementation and valid-row execution call this method. - /// - /// Reuse `values` when it already has the required physical representation. Any new payload - /// buffers must use `allocator`. - fn build(values: BufferMut, allocator: &BufferAllocatorRef) -> ArrayRef; + /// Finishing the buffer must produce an all-valid column whose dtype matches + /// [`element_dtype`](Self::element_dtype) except for outer nullability. + fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer; /// Map a contiguous row source directly into an all-valid column. /// - /// The default collects into a [`BufferMut`] using `allocator`, then calls [`build`](Self::build). + /// The default writes into the storage returned by [`allocate`](Self::allocate), then finishes it. /// An output type can override this method when its physical representation supports a more /// efficient bulk mapping. The implementation **must** call `apply` exactly once for every /// source row in increasing order and return the same values as the default implementation. @@ -57,14 +57,40 @@ pub trait OutputElement: 'static + Sized + Default { F: Fn(S::Item) -> Self, { let row_count = source.len(); - let mut values = allocator.with_capacity::(row_count); - let output = &mut values.spare_capacity_mut()[..row_count]; + let mut values = Self::allocate(row_count, allocator); + let output = &mut values.slots()[..row_count]; source.map_into(output, apply); // SAFETY: normal completion of `map_into` initializes every output slot exactly once. - unsafe { values.set_len(row_count) }; - - Self::build(values, allocator) + unsafe { values.finish(row_count, allocator) } } } + +/// Engine-owned storage for collecting independent row values. +/// +/// The executor initializes a prefix of [`slots`](Self::slots), then calls [`finish`](Self::finish). +/// Dropping the buffer instead abandons the output, including after an error or unwind. +/// +/// # Safety +/// +/// - Access through this trait must preserve slot count and contents between calls to +/// [`slots`](Self::slots), including across moves of the buffer. +/// - Dropping the buffer must be safe with any subset of its slots initialized. +/// +/// Violating these requirements can cause undefined behavior in an executor or sink. +pub unsafe trait OutputBuffer: Sized { + /// Writable capacity whose initialized prefix is published by [`finish`](Self::finish). + fn slots(&mut self) -> &mut [MaybeUninit]; + + /// Publish the first `len` slots as an all-valid column, reusing storage where possible. + /// + /// The column must contain exactly `len` rows in slot order. Any new payload buffers must + /// use `allocator`. + /// + /// # Safety + /// + /// The first `len` slots must exist and contain initialized values. Violating this requirement + /// can cause undefined behavior. + unsafe fn finish(self, len: usize, allocator: &BufferAllocatorRef) -> ArrayRef; +} 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 971b1add143..dbe25b2ae89 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 @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::MaybeUninit; + use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; @@ -18,6 +20,7 @@ use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::scalar::ScalarValue; use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::validity::Validity; @@ -100,11 +103,29 @@ unsafe impl InputElement for T { } impl OutputElement for T { + type Buffer = BufferMut; + fn element_dtype() -> DType { DType::Primitive(T::PTYPE, Nullability::NonNullable) } - fn build(values: BufferMut, _allocator: &BufferAllocatorRef) -> ArrayRef { - PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array() + fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { + allocator.with_capacity(rows) + } +} + +// SAFETY: clearing the length preserves the contents and exposes the same allocation each time. +// Native values require no destruction when a partially initialized buffer is abandoned. +unsafe impl OutputBuffer for BufferMut { + fn slots(&mut self) -> &mut [MaybeUninit] { + self.clear(); + self.spare_capacity_mut() + } + + unsafe fn finish(mut self, len: usize, _allocator: &BufferAllocatorRef) -> ArrayRef { + // SAFETY: the caller initialized the first `len` slots of this buffer's spare capacity. + unsafe { self.set_len(len) }; + + PrimitiveArray::new(self.freeze(), Validity::NonNullable).into_array() } } 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 a146a35ffcb..0f5037d0570 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -16,6 +16,7 @@ mod element; pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; +pub use element::OutputBuffer; pub use element::OutputElement; pub use element::Utf8Column; pub use element::Utf8View; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs index e266e0e986d..d474177d9a4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs @@ -10,7 +10,6 @@ use std::mem::MaybeUninit; use std::sync::Arc; use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -22,6 +21,7 @@ use crate::arrays::FixedSizeListArray; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar_fn::unstable::row::FillDefault; +use crate::scalar_fn::unstable::row::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::ViewLen; use crate::validity::Validity; @@ -87,9 +87,9 @@ impl FillDefault for FixedSizeRows<'_, T> { /// and validates that physical parameter before calling [`RowVisitor::visit_into`]. /// /// [`RowVisitor::visit_into`]: crate::scalar_fn::unstable::row::RowVisitor::visit_into -pub struct FixedSizeListSink { +pub struct FixedSizeListSink { /// Spare flat storage written one fixed-size row at a time. - values: BufferMut, + values: T::Buffer, /// Allocator for any physical conversion when the sink finishes. allocator: BufferAllocatorRef, /// The number of elements in each output row. @@ -102,7 +102,8 @@ pub struct FixedSizeListSink { // shape for its lifetime. Each row is one disjoint `width`-element slice. `InitializedRow::fill` // requires the entire current row and preservation of its initialization until the callback // returns its private token. `FixedSizeRows::fill_default` writes every flat element before masked -// traversal. `values` retains length zero until every row is safe to publish in `finish`. +// traversal. `OutputBuffer` preserves initialized slots across row views and permits abandoning +// partially initialized storage. unsafe impl OutputSink for FixedSizeListSink { type Params = usize; type Rows<'a> = FixedSizeRows<'a, T>; @@ -131,7 +132,7 @@ unsafe impl OutputSink for FixedSizeListSink< })?; Ok(Self { - values: allocator.with_capacity(element_capacity), + values: T::allocate(element_capacity, allocator), allocator: allocator.clone(), width, row_count: rows, @@ -140,7 +141,7 @@ unsafe impl OutputSink for FixedSizeListSink< fn rows(&mut self) -> Self::Rows<'_> { FixedSizeRows { - elements: &mut self.values.spare_capacity_mut()[..self.row_count * self.width], + elements: &mut self.values.slots()[..self.row_count * self.width], width: self.width, row_count: self.row_count, } @@ -155,14 +156,11 @@ unsafe impl OutputSink for FixedSizeListSink< unsafe { rows.elements.get_unchecked_mut(start..end) } } - unsafe fn finish(mut self) -> VortexResult { + unsafe fn finish(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, &self.allocator); + // SAFETY: the caller guarantees every row's `width` elements were initialized. + let elements = unsafe { self.values.finish(element_count, &self.allocator) }; let lists = FixedSizeListArray::new( elements, fixed_size_list_size(self.width), diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs index 0f5bbac9e73..785dda8c573 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs @@ -9,13 +9,13 @@ use std::mem::MaybeUninit; use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; use super::OutputSink; use crate::ArrayRef; use crate::dtype::DType; use crate::scalar_fn::unstable::row::FillDefault; +use crate::scalar_fn::unstable::row::OutputBuffer; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::ViewLen; @@ -58,11 +58,11 @@ impl InitializedElement { /// 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 { +/// safe because [`OutputBuffer`] permits abandoning partially initialized storage. The `T: Copy` +/// bound means that initialized elements require no destruction. +pub struct UninitElementSink { /// Spare storage written in increasing row order. - values: BufferMut, + values: T::Buffer, /// Allocator for any physical conversion when the sink finishes. allocator: BufferAllocatorRef, /// The number of slots exposed to the row loop and initialized before finishing. @@ -94,7 +94,8 @@ impl FillDefault for UninitElementRows<'_, T> { // 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 default // skipped-row initializer fills every slot with `T::default()` through -// `UninitElementRows::fill_default` before masked traversal. +// `UninitElementRows::fill_default` before masked traversal. `OutputBuffer` preserves initialized +// slots across row views and permits abandoning partially initialized storage. unsafe impl OutputSink for UninitElementSink { type Params = (); type Rows<'a> = UninitElementRows<'a, T>; @@ -111,14 +112,14 @@ unsafe impl OutputSink for UninitElementSink< allocator: &BufferAllocatorRef, ) -> VortexResult { Ok(Self { - values: allocator.with_capacity(rows), + values: T::allocate(rows, allocator), allocator: allocator.clone(), row_count: rows, }) } fn rows(&mut self) -> Self::Rows<'_> { - UninitElementRows(&mut self.values.spare_capacity_mut()[..self.row_count]) + UninitElementRows(&mut self.values.slots()[..self.row_count]) } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { @@ -126,11 +127,8 @@ unsafe impl OutputSink for UninitElementSink< unsafe { rows.0.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, &self.allocator)) + unsafe fn finish(self) -> VortexResult { + // SAFETY: the caller guarantees every exposed slot was initialized. + Ok(unsafe { self.values.finish(self.row_count, &self.allocator) }) } } From b93f38142a1603371d3f5ac49771b852ce9e66d9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 17:15:57 -0400 Subject: [PATCH 3/6] Fix allocator test helper lints Signed-off-by: Connor Tsui --- vortex-array/src/memory.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/memory.rs b/vortex-array/src/memory.rs index 7cac795da2c..02d79b4bcb1 100644 --- a/vortex-array/src/memory.rs +++ b/vortex-array/src/memory.rs @@ -79,13 +79,13 @@ pub(crate) mod test_allocator { use std::alloc::Layout; use std::ptr::NonNull; use std::sync::Arc; - use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; use allocator_api2::alloc::Global; + use parking_lot::Mutex; use super::BufferAllocatorRef; @@ -118,11 +118,10 @@ pub(crate) mod test_allocator { "allocation ownership needs a nonempty payload", ); let start = values.as_ptr() as usize; - let end = start + std::mem::size_of_val(values); + let end = start + size_of_val(values); assert!( self.allocations .lock() - .unwrap() .iter() .any(|&(base, size)| base <= start && end <= base + size), "the returned payload must be backed by the configured allocator", @@ -130,7 +129,7 @@ pub(crate) mod test_allocator { } pub(crate) fn live_allocations(&self) -> usize { - self.allocations.lock().unwrap().len() + self.allocations.lock().len() } } @@ -145,7 +144,6 @@ pub(crate) mod test_allocator { self.0 .allocations .lock() - .unwrap() .push((allocation.cast::().as_ptr() as usize, allocation.len())); Ok(allocation) } @@ -154,7 +152,6 @@ pub(crate) mod test_allocator { self.0 .allocations .lock() - .unwrap() .retain(|&(base, _)| base != ptr.as_ptr() as usize); // SAFETY: this allocation came from Global with the supplied layout. unsafe { Global.deallocate(ptr, layout) } From 0a0feb7af56d7fc9151832a46235b0a58d239375 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 17:17:35 -0400 Subject: [PATCH 4/6] Match failure evidence width in zero-sized output tests Signed-off-by: Connor Tsui --- .../scalar_fn/unstable/row/execute/tests.rs | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index 96ee66bdb98..35a8b102e5c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -4,6 +4,7 @@ //! Verifies ownership of RowFn output allocations independently of input decoding. use std::mem::MaybeUninit; +use std::ops::BitOrAssign; use rstest::rstest; use vortex_buffer::BufferAllocatorRef; @@ -42,6 +43,7 @@ use crate::dtype::DType; use crate::memory::MemorySessionExt; use crate::memory::test_allocator::tracking_allocator; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::FixedSizeListSink; use crate::scalar_fn::unstable::row::InitializedElement; use crate::scalar_fn::unstable::row::OutputBuffer; @@ -61,7 +63,7 @@ enum Traversal { Filtered, } -fn collect_owned( +fn collect_owned( traversal: Traversal, args: &VecExecutionArgs, valid: &MaskValuesRef, @@ -75,19 +77,19 @@ fn collect_owned( |_| (), |_, (value,)| apply(value), ), - Traversal::Fallible => execute_owned::<(i64,), Out, (), bool>( + Traversal::Fallible => execute_owned::<(i64,), Out, (), Fail>( args, ctx, |_| (), - |_, (value,)| (apply(value), false), + |_, (value,)| (apply(value), Fail::default()), |_| Ok(()), ), Traversal::DenseAttempt => { - match execute_owned_dense_attempt::<(i64,), Out, (), bool>( + match execute_owned_dense_attempt::<(i64,), Out, (), Fail>( args, ctx, |_| (), - |_, (value,)| (apply(value), false), + |_, (value,)| (apply(value), Fail::default()), |_| Ok(()), )? { DenseAttempt::Values(values) => Ok(values), @@ -156,9 +158,9 @@ fn owned_payload_uses_context_allocator( .with_allocator(allocator); let output = if boolean { - collect_owned(traversal, &args, &valid, &mut ctx, |value| value % 2 == 0)? + collect_owned::<_, bool>(traversal, &args, &valid, &mut ctx, |value| value % 2 == 0)? } else { - collect_owned(traversal, &args, &valid, &mut ctx, |value| value)? + collect_owned::<_, bool>(traversal, &args, &valid, &mut ctx, |value| value)? }; if boolean { tracker.assert_owns(output.as_::().to_bit_buffer().inner().as_slice()); @@ -406,6 +408,14 @@ fn fixed_size_list_payload_uses_allocator() -> VortexResult<()> { Ok(()) } +// Zero-sized outputs require failure evidence that is also zero-sized. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + /// A zero-sized element whose collection storage does not depend on Vortex buffers. #[derive(Clone, Copy, Default)] struct One; @@ -450,7 +460,7 @@ fn zero_sized_output_uses_its_own_storage(#[case] traversal: Traversal) -> Vorte .create_execution_ctx() .with_allocator(allocator); - let output = collect_owned::(traversal, &args, &valid, &mut ctx, |_| One)?; + let output = collect_owned::(traversal, &args, &valid, &mut ctx, |_| One)?; assert_arrays_eq!(&output, &expected, &mut ctx); tracker.assert_owns(output.as_::().as_slice::()); From 5350a60019aea81ee47022e2affbacb6849ae008 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 17:52:24 -0400 Subject: [PATCH 5/6] Use the execution allocator for the new Boolean dense retry path Signed-off-by: Connor Tsui --- .../unstable/row/execute/packed_bool.rs | 4 +- .../scalar_fn/unstable/row/execute/tests.rs | 40 ++++++++++++++++--- 2 files changed, 36 insertions(+), 8 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs index d494558d0ff..5fb3d0cadc1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs @@ -162,9 +162,9 @@ where }; let values = if MULTIVERSIONED { - BitBuffer::collect_bool_multiversioned(row_count, collect) + BitBuffer::collect_bool_multiversioned_in(row_count, collect, ctx.allocator().clone()) } else { - BitBuffer::collect_bool(row_count, collect) + BitBuffer::collect_bool_in(row_count, collect, ctx.allocator().clone()) }; match finish_failure(state.failure) { diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index 35a8b102e5c..24df6efc976 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -15,6 +15,7 @@ use vortex_mask::Mask; use vortex_mask::MaskValuesRef; use super::DenseAttempt; +use super::execute_bool_dense_attempt; use super::execute_owned; use super::execute_owned_bool; use super::execute_owned_dense_attempt; @@ -175,9 +176,12 @@ fn owned_payload_uses_context_allocator( } #[rstest] +#[case::infallible(Traversal::Infallible)] +#[case::fallible(Traversal::Fallible)] +#[case::dense_attempt(Traversal::DenseAttempt)] fn packed_boolean_payload_uses_allocator( + #[case] traversal: Traversal, #[values(false, true)] multiversioned: bool, - #[values(false, true)] deferred: bool, ) -> VortexResult<()> { let args = canonical_args(Traversal::Infallible, false); let (allocator, tracker) = tracking_allocator(); @@ -185,27 +189,51 @@ fn packed_boolean_payload_uses_allocator( .create_execution_ctx() .with_allocator(allocator); - let output = match (multiversioned, deferred) { - (false, false) => { + let output = match (multiversioned, traversal) { + (false, Traversal::Infallible) => { execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 2)? } - (true, false) => { + (true, Traversal::Infallible) => { execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, |(value,)| value > 2)? } - (false, true) => execute_owned_bool::<(i64,), (), bool, false>( + (false, Traversal::Fallible) => execute_owned_bool::<(i64,), (), bool, false>( &args, &mut ctx, |_| (), |_, (value,)| (value > 2, false), |_| Ok(()), )?, - (true, true) => execute_owned_bool::<(i64,), (), bool, true>( + (true, Traversal::Fallible) => execute_owned_bool::<(i64,), (), bool, true>( &args, &mut ctx, |_| (), |_, (value,)| (value > 2, false), |_| Ok(()), )?, + (multiversioned, Traversal::DenseAttempt) => { + let attempt = if multiversioned { + execute_bool_dense_attempt::<(i64,), (), bool, true>( + &args, + &mut ctx, + |_| (), + |_, (value,)| (value > 2, false), + |_| Ok(()), + )? + } else { + execute_bool_dense_attempt::<(i64,), (), bool, false>( + &args, + &mut ctx, + |_| (), + |_, (value,)| (value > 2, false), + |_| Ok(()), + )? + }; + match attempt { + DenseAttempt::Values(values) => values, + DenseAttempt::DeferredError(error) => return Err(error), + } + } + _ => vortex_bail!("this test traversal requires packed Boolean output"), }; tracker.assert_owns(output.as_::().to_bit_buffer().inner().as_slice()); assert_eq!(tracker.live_allocations(), 1); From 8e3b4c08f20bb0b83fc7712341bbcaee59f431b7 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 24 Sep 2026 15:09:13 -0400 Subject: [PATCH 6/6] refactor(array): rename output allocation to with_capacity Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/unstable/row/batch/tests.rs | 2 +- vortex-array/src/scalar_fn/unstable/row/execute/owned.rs | 6 +++--- vortex-array/src/scalar_fn/unstable/row/execute/retry.rs | 2 +- vortex-array/src/scalar_fn/unstable/row/execute/tests.rs | 4 ++-- .../src/scalar_fn/unstable/row/types/element/bool.rs | 2 +- .../src/scalar_fn/unstable/row/types/element/output.rs | 7 ++++--- .../src/scalar_fn/unstable/row/types/element/primitive.rs | 2 +- .../scalar_fn/unstable/row/types/sink/fixed_size_list.rs | 2 +- .../scalar_fn/unstable/row/types/sink/uninit_element.rs | 2 +- 9 files changed, 15 insertions(+), 14 deletions(-) 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 4baae4df3d3..d2ae3bdf93f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -261,7 +261,7 @@ impl OutputElement for NullProducingI64 { DType::from(i64::PTYPE) } - fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { + fn with_capacity(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { allocator.with_capacity(rows) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index b34d89a391e..8a0b826a461 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -142,7 +142,7 @@ where let prepared = prepare(Args::const_values(&columns)); let valid_rows = valid.bit_buffer(); - let mut values = Out::allocate(valid_rows.len(), ctx.allocator()); + let mut values = Out::with_capacity(valid_rows.len(), ctx.allocator()); let output = &mut values.slots()[..valid_rows.len()]; for slot in output.iter_mut() { @@ -221,7 +221,7 @@ where ); let prepared = prepare(Args::const_values(&columns)); - let mut values = Out::allocate(row_count, ctx.allocator()); + let mut values = Out::with_capacity(row_count, ctx.allocator()); let output = &mut values.slots()[..row_count]; for slot in output.iter_mut() { @@ -288,7 +288,7 @@ where let prepared = prepare(Args::const_values(&columns)); let row_count = args.row_count(); - let mut values = Out::allocate(row_count, ctx.allocator()); + let mut values = Out::with_capacity(row_count, ctx.allocator()); let output = &mut values.slots()[..row_count]; let Some(source) = decoded_source::(&columns, row_count) else { diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs index 8ca4babcff3..ebf67461e6c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/retry.rs @@ -62,7 +62,7 @@ where let prepared = prepare(Args::const_values(&columns)); let row_count = args.row_count(); - let mut values = Out::allocate(row_count, ctx.allocator()); + let mut values = Out::with_capacity(row_count, ctx.allocator()); let output = &mut values.slots()[..row_count]; let failure_evidence = if let Some(views) = Args::views_if_no_consts(&columns) { diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index 24df6efc976..80d954bd8db 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -321,7 +321,7 @@ fn sink_payloads_use_allocator(#[case] traversal: Traversal) -> VortexResult<()> #[test] fn primitive_finish_reuses_allocation() { let (allocator, tracker) = tracking_allocator(); - let mut values = i64::allocate(3, &allocator); + let mut values = i64::with_capacity(3, &allocator); let slots = &mut values.slots()[..3]; let ptr = slots.as_ptr().cast::(); @@ -455,7 +455,7 @@ impl OutputElement for One { i64::element_dtype() } - fn allocate(rows: usize, _allocator: &BufferAllocatorRef) -> Self::Buffer { + fn with_capacity(rows: usize, _allocator: &BufferAllocatorRef) -> Self::Buffer { vec![MaybeUninit::uninit(); rows] } } 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 707b6e18e94..9b517ef89d9 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 @@ -101,7 +101,7 @@ impl OutputElement for bool { DType::Bool(Nullability::NonNullable) } - fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { + fn with_capacity(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { allocator.with_capacity(rows) } 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 3840bec8b05..e2d944d9516 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 @@ -36,11 +36,12 @@ pub trait OutputElement: 'static + Sized + Default { /// /// Finishing the buffer must produce an all-valid column whose dtype matches /// [`element_dtype`](Self::element_dtype) except for outer nullability. - fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer; + fn with_capacity(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer; /// Map a contiguous row source directly into an all-valid column. /// - /// The default writes into the storage returned by [`allocate`](Self::allocate), then finishes it. + /// The default writes into the storage returned by [`with_capacity`](Self::with_capacity), then + /// finishes it. /// An output type can override this method when its physical representation supports a more /// efficient bulk mapping. The implementation **must** call `apply` exactly once for every /// source row in increasing order and return the same values as the default implementation. @@ -57,7 +58,7 @@ pub trait OutputElement: 'static + Sized + Default { F: Fn(S::Item) -> Self, { let row_count = source.len(); - let mut values = Self::allocate(row_count, allocator); + let mut values = Self::with_capacity(row_count, allocator); let output = &mut values.slots()[..row_count]; source.map_into(output, apply); 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 dbe25b2ae89..5f1e3fd2dd9 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 @@ -109,7 +109,7 @@ impl OutputElement for T { DType::Primitive(T::PTYPE, Nullability::NonNullable) } - fn allocate(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { + fn with_capacity(rows: usize, allocator: &BufferAllocatorRef) -> Self::Buffer { allocator.with_capacity(rows) } } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs index d474177d9a4..5a2f86580e3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/fixed_size_list.rs @@ -132,7 +132,7 @@ unsafe impl OutputSink for FixedSizeListSink< })?; Ok(Self { - values: T::allocate(element_capacity, allocator), + values: T::with_capacity(element_capacity, allocator), allocator: allocator.clone(), width, row_count: rows, diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs index 785dda8c573..700bfb595e4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink/uninit_element.rs @@ -112,7 +112,7 @@ unsafe impl OutputSink for UninitElementSink< allocator: &BufferAllocatorRef, ) -> VortexResult { Ok(Self { - values: T::allocate(rows, allocator), + values: T::with_capacity(rows, allocator), allocator: allocator.clone(), row_count: rows, })