From 8c6e97cbddb3008cafe03bb8504b3f13e18f81a4 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 27 Aug 2026 12:41:31 -0400 Subject: [PATCH 1/3] refactor(array): use buffer allocator in execution context Signed-off-by: Nicholas Gates --- Cargo.lock | 2 + vortex-array/Cargo.toml | 1 + .../src/arrays/chunked/vtable/canonical.rs | 57 ++- vortex-array/src/builders/mod.rs | 6 +- vortex-array/src/executor.rs | 50 ++- vortex-array/src/memory.rs | 343 ++---------------- .../src/scalar_fn/fns/variant_get/mod.rs | 3 +- vortex-buffer/src/allocation.rs | 9 + vortex-file/Cargo.toml | 1 + vortex-file/src/open.rs | 23 +- vortex-io/src/object_store/read_at.rs | 26 +- vortex-io/src/std_file/read_at.rs | 15 +- 12 files changed, 151 insertions(+), 385 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0b59cb19243..c42bc8af47a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10491,6 +10491,7 @@ dependencies = [ name = "vortex-array" version = "0.1.0" dependencies = [ + "allocator-api2", "arbitrary", "arc-swap", "arcref", @@ -11042,6 +11043,7 @@ dependencies = [ name = "vortex-file" version = "0.1.0" dependencies = [ + "allocator-api2", "async-stream", "async-trait", "bytes", diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index fb9ea651ad0..f6b06544baf 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -83,6 +83,7 @@ serde = ["dep:serde", "vortex-buffer/serde", "vortex-mask/serde"] unstable_row_fns = [] [dev-dependencies] +allocator-api2 = { workspace = true } divan = { workspace = true } futures = { workspace = true, features = ["executor"] } insta = { workspace = true } diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 8cc8c7aa4a2..df6e5bd3739 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; -use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -30,7 +29,6 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; -use crate::memory::HostAllocatorExt; use crate::validity::Validity; pub(super) fn _canonicalize( @@ -71,7 +69,8 @@ pub(super) fn _canonicalize( } DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?), _ => { - let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); + let mut builder = + builder_with_capacity_in(ctx.allocator().clone(), array.dtype(), array.len()); array.array().append_to_builder(builder.as_mut(), ctx)?; builder.finish_into_canonical(ctx) } @@ -180,10 +179,10 @@ fn swizzle_list_chunks( // We (somewhat arbitrarily) choose `u64` for our offsets and sizes here. These can always be // narrowed later by the compressor. let allocator = ctx.allocator(); - let mut offsets = allocator.allocate_typed::(len)?; - let mut sizes = allocator.allocate_typed::(len)?; - let offsets_out: &mut [u64] = offsets.as_mut_slice_typed::()?; - let sizes_slice_out: &mut [u64] = sizes.as_mut_slice_typed::()?; + let mut offsets = allocator.zeroed::(len); + let mut sizes = allocator.zeroed::(len); + let offsets_out = offsets.as_mut_slice(); + let sizes_slice_out = sizes.as_mut_slice(); let mut next_list = 0usize; for chunk in chunks { @@ -229,16 +228,8 @@ fn swizzle_list_chunks( unsafe { ChunkedArray::new_unchecked(list_elements_chunks, elem_dtype.clone()) } .into_array(); - let offsets = PrimitiveArray::new( - Buffer::::from_byte_buffer(offsets.freeze()), - Validity::NonNullable, - ) - .into_array(); - let sizes = PrimitiveArray::new( - Buffer::::from_byte_buffer(sizes.freeze()), - Validity::NonNullable, - ) - .into_array(); + let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable).into_array(); + let sizes = PrimitiveArray::new(sizes.freeze(), Validity::NonNullable).into_array(); // SAFETY: // - `offsets` and `sizes` are non-nullable u64 arrays of the same length @@ -286,11 +277,17 @@ fn swizzle_fixed_size_list_chunks( #[cfg(test)] mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; use std::sync::Arc; use std::sync::LazyLock; 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 vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -318,10 +315,7 @@ mod tests { use crate::dtype::DType::Variant as VariantDType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; - use crate::memory::DefaultHostAllocator; - use crate::memory::HostAllocator; use crate::memory::MemorySessionExt; - use crate::memory::WritableHostBuffer; use crate::scalar::Scalar; use crate::validity::Validity; @@ -333,14 +327,16 @@ mod tests { allocations: Arc, } - impl HostAllocator for CountingAllocator { - fn allocate( - &self, - len: usize, - alignment: vortex_buffer::Alignment, - ) -> VortexResult { + // SAFETY: this forwards memory operations to Global and only counts allocations. + unsafe impl Allocator for CountingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { self.allocations.fetch_add(1, Ordering::Relaxed); - DefaultHostAllocator.allocate(len, alignment) + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: ptr and layout came from Global. + unsafe { Global.deallocate(ptr, layout) } } } @@ -663,9 +659,10 @@ mod tests { #[test] fn list_canonicalize_uses_memory_session_allocator() { let allocations = Arc::new(AtomicUsize::new(0)); - let session = crate::array_session().with_allocator(Arc::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); + let session = + crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator { + allocations: Arc::clone(&allocations), + })); let mut ctx = session.create_execution_ctx(); let l1 = ListArray::try_new( diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 22c806f74a1..ea8ee4d74a0 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -47,7 +47,7 @@ use crate::canonical::Canonical; use crate::dtype::DType; use crate::match_each_decimal_value_type; use crate::match_each_native_ptype; -use crate::memory::HostAllocatorRef; +use crate::memory::BufferAllocatorRef; use crate::scalar::Scalar; mod lazy_null_builder; @@ -447,9 +447,9 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box Box { diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 4e87ff5e7dd..04c4a31ce4a 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -12,6 +12,7 @@ //! See for the full execution //! narrative, diagrams, and walkthroughs. +use std::cell::OnceCell; use std::env::VarError; use std::fmt; use std::fmt::Display; @@ -38,7 +39,7 @@ use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::matcher::Matcher; -use crate::memory::HostAllocatorRef; +use crate::memory::BufferAllocatorRef; use crate::memory::MemorySessionExt; use crate::optimizer::ArrayOptimizer; use crate::optimizer::kernels::ArrayKernelsExt; @@ -292,7 +293,7 @@ impl ArrayRef { if current_builder.is_none() { trace_op!(record_builder_start(&array)); current_builder = Some(builder_with_capacity_in( - ctx.allocator(), + ctx.allocator().clone(), array.dtype(), array.len(), )); @@ -350,6 +351,8 @@ struct StackFrame { #[derive(Debug, Clone)] pub struct ExecutionCtx { session: VortexSession, + // OnceCell avoids cloning the session allocator when a context does not allocate. + allocator: OnceCell, execute_parent_kernels: Arc, #[cfg(debug_assertions)] id: usize, @@ -367,6 +370,7 @@ impl ExecutionCtx { let execute_parent_kernels = session.kernels().execute_parent_snapshot(); Self { session, + allocator: OnceCell::new(), execute_parent_kernels, #[cfg(debug_assertions)] id: { @@ -383,9 +387,15 @@ impl ExecutionCtx { &self.session } - /// Get the session-scoped host allocator for this execution context. - pub fn allocator(&self) -> HostAllocatorRef { - self.session.allocator() + /// Get the allocator for this execution context. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.allocator.get_or_init(|| self.session.allocator()) + } + + /// Set the allocator for this execution context. + pub fn with_allocator(mut self, allocator: BufferAllocatorRef) -> Self { + self.allocator = OnceCell::from(allocator); + self } /// Log an execution step at the current depth. @@ -531,7 +541,8 @@ impl Executable for ArrayRef { ExecutionStep::AppendChild(_) => { // Single-step: build the entire parent via the builder path. trace_op!(record_builder_start(&array)); - let builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len()); + let builder = + builder_with_capacity_in(ctx.allocator().clone(), array.dtype(), array.len()); let mut builder = execute_into_builder(array, builder, ctx)?; let output = builder.finish(); trace_op!(record_builder_finish(&output)); @@ -926,6 +937,7 @@ impl VortexSessionExecute for VortexSession { #[cfg(test)] mod tests { + use vortex_session::SessionExt; use vortex_session::VortexSession; use super::*; @@ -933,6 +945,9 @@ mod tests { use crate::VortexSessionExecute; use crate::arrays::Bool; use crate::arrays::Primitive; + use crate::memory::BufferAllocatorRef; + use crate::memory::MemorySession; + use crate::memory::MemorySessionExt; use crate::optimizer::kernels::ExecuteParentFn; use crate::optimizer::kernels::KernelSession; use crate::optimizer::kernels::execute_parent_key; @@ -974,4 +989,27 @@ mod tests { let after_registration = session.create_execution_ctx(); assert!(after_registration.execute_parent_kernels.contains_key(&key)); } + + #[test] + fn execution_ctx_allocator_override() { + let first = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); + let second = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); + let third = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator); + let session = VortexSession::empty() + .with::() + .with_allocator(first.clone()); + let ctx = session.create_execution_ctx(); + + session + .get_mut::() + .set_allocator(third.clone()); + + assert!(session.allocator().ptr_eq(&third)); + assert!(ctx.allocator().ptr_eq(&third)); + + let ctx = ctx.with_allocator(second.clone()); + session.get_mut::().set_allocator(first); + + assert!(ctx.allocator().ptr_eq(&second)); + } } diff --git a/vortex-array/src/memory.rs b/vortex-array/src/memory.rs index b52b3442136..3d328915933 100644 --- a/vortex-array/src/memory.rs +++ b/vortex-array/src/memory.rs @@ -1,203 +1,44 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Session-scoped memory allocation for host-side buffers. +//! Session-scoped buffer allocation. use std::any::Any; -use std::fmt::Debug; -use std::mem::size_of; -use std::sync::Arc; -use bytes::Bytes; -use vortex_buffer::Alignment; -use vortex_buffer::Buffer; -use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; +pub use vortex_buffer::BufferAllocator; +pub use vortex_buffer::BufferAllocatorRef; +pub use vortex_buffer::StaticBufferAllocator; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::VortexSession; -/// Mutable host buffer contract used by [`WritableHostBuffer`]. -pub trait HostBufferMut: Send + 'static { - /// Returns the logical byte length of the buffer. - fn len(&self) -> usize; - - /// Whether the buffer is empty. - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Returns the alignment of the buffer. - fn alignment(&self) -> Alignment; - - /// Returns mutable access to the writable byte range. - fn as_mut_slice(&mut self) -> &mut [u8]; - - /// Freeze the buffer into an immutable [`ByteBuffer`]. - fn freeze(self: Box) -> ByteBuffer; -} - -/// Exact-size writable host buffer returned by a [`HostAllocator`]. -pub struct WritableHostBuffer { - inner: Box, -} - -impl WritableHostBuffer { - /// Create a writable host buffer from an implementation of [`HostBufferMut`]. - pub fn new(inner: Box) -> Self { - Self { inner } - } - - /// Returns the logical byte length of the buffer. - pub fn len(&self) -> usize { - self.inner.len() - } - - /// Returns true when the buffer has zero bytes. - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Returns the alignment of the buffer. - pub fn alignment(&self) -> Alignment { - self.inner.alignment() - } - - /// Returns mutable access to the writable byte range. - pub fn as_mut_slice(&mut self) -> &mut [u8] { - self.inner.as_mut_slice() - } - - /// Returns mutable access to the buffer as a typed slice. - pub fn as_mut_slice_typed(&mut self) -> VortexResult<&mut [T]> { - vortex_ensure!( - size_of::() != 0, - InvalidArgument: "Cannot create typed mutable slice for zero-sized type {}", - std::any::type_name::() - ); - vortex_ensure!( - self.alignment().is_aligned_to(Alignment::of::()), - InvalidArgument: "Buffer is not sufficiently aligned for type {}", - std::any::type_name::() - ); - - let bytes = self.as_mut_slice(); - let byte_len = bytes.len(); - let ptr = bytes.as_mut_ptr(); - let type_size = size_of::(); - - vortex_ensure!( - byte_len.is_multiple_of(type_size), - InvalidArgument: "Buffer length {byte_len} is not a multiple of {} for {}", - type_size, - std::any::type_name::() - ); - - // SAFETY: We checked size divisibility and pointer alignment for `T`, - // and we have exclusive mutable access to the underlying bytes. - Ok(unsafe { std::slice::from_raw_parts_mut(ptr.cast::(), byte_len / type_size) }) - } - - /// Freeze the writable buffer into an immutable [`ByteBuffer`]. - pub fn freeze(self) -> ByteBuffer { - self.inner.freeze() - } - - /// Freeze the writable buffer into a typed immutable [`Buffer`]. - pub fn freeze_typed(self) -> VortexResult> { - vortex_ensure!( - size_of::() != 0, - InvalidArgument: "Cannot freeze typed buffer for zero-sized type {}", - std::any::type_name::() - ); - - let buffer = self.freeze(); - let byte_len = buffer.len(); - let type_size = size_of::(); - let type_align = Alignment::of::(); - - vortex_ensure!( - byte_len.is_multiple_of(type_size), - InvalidArgument: "Buffer length {byte_len} is not a multiple of {} for {}", - type_size, - std::any::type_name::() - ); - vortex_ensure!( - buffer.is_aligned(type_align), - InvalidArgument: "Buffer pointer is not aligned to {} for {}", - type_align, - std::any::type_name::() - ); - - Ok(Buffer::from_byte_buffer(buffer)) - } -} - -impl Debug for WritableHostBuffer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("WritableHostBuffer") - .field("len", &self.len()) - .field("alignment", &self.alignment()) - .finish() - } -} - -/// Allocator for exact-size writable host buffers. -pub trait HostAllocator: Debug + Send + Sync + 'static { - /// Allocate a writable host buffer with the requested byte length and alignment. - fn allocate(&self, len: usize, alignment: Alignment) -> VortexResult; -} - -/// Shared allocator reference used throughout session-scoped memory APIs. -pub type HostAllocatorRef = Arc; - -/// Extension methods for [`HostAllocator`]s. -pub trait HostAllocatorExt: HostAllocator { - /// Allocate host memory for `len` elements of `T` using `Alignment::of::()`. - fn allocate_typed(&self, len: usize) -> VortexResult { - let bytes = len.checked_mul(size_of::()).ok_or_else(|| { - vortex_err!( - "Typed host allocation overflow for type {} and len {}", - std::any::type_name::(), - len - ) - })?; - self.allocate(bytes, Alignment::of::()) - } -} - -impl HostAllocatorExt for A {} - /// Session-scoped memory configuration for Vortex arrays. #[derive(Clone, Debug)] pub struct MemorySession { - allocator: HostAllocatorRef, + allocator: BufferAllocatorRef, } impl MemorySession { - /// Creates a new session memory configuration using the provided allocator. - pub fn new(allocator: HostAllocatorRef) -> Self { + /// Creates a new memory configuration using the provided allocator. + pub fn new(allocator: BufferAllocatorRef) -> Self { Self { allocator } } /// Returns the configured allocator. - pub fn allocator(&self) -> HostAllocatorRef { - Arc::clone(&self.allocator) + pub fn allocator(&self) -> BufferAllocatorRef { + self.allocator.clone() } /// Updates the configured allocator. - pub fn set_allocator(&mut self, allocator: HostAllocatorRef) { + pub fn set_allocator(&mut self, allocator: BufferAllocatorRef) { self.allocator = allocator; } } impl Default for MemorySession { fn default() -> Self { - Self::new(Arc::new(DefaultHostAllocator)) + Self::new(BufferAllocatorRef::statically_allocated()) } } @@ -211,21 +52,20 @@ impl SessionVar for MemorySession { } } -/// Extension trait for accessing session-scoped memory configuration. +/// Extension methods for session-scoped buffer allocation. pub trait MemorySessionExt: SessionExt { - /// Returns the memory session for this execution/session context. + /// Returns the memory configuration. fn memory(&self) -> SessionGuard<'_, MemorySession> { self.get::() } - /// Returns the configured host allocator for this execution/session context. - fn allocator(&self) -> HostAllocatorRef { + /// Returns the configured buffer allocator. + fn allocator(&self) -> BufferAllocatorRef { self.memory().allocator() } - /// Configures the session to use `allocator` as its host allocator, mutating it in place and - /// returning it for chaining. - fn with_allocator(self, allocator: HostAllocatorRef) -> VortexSession { + /// Configures the session allocator and returns the session. + fn with_allocator(self, allocator: BufferAllocatorRef) -> VortexSession { let session = self.session(); session.get_mut::().set_allocator(allocator); session @@ -234,157 +74,18 @@ pub trait MemorySessionExt: SessionExt { impl MemorySessionExt for S {} -/// Default host allocator. -#[derive(Debug, Default)] -pub struct DefaultHostAllocator; - -impl HostAllocator for DefaultHostAllocator { - fn allocate(&self, len: usize, alignment: Alignment) -> VortexResult { - let mut buffer = ByteBufferMut::with_capacity_aligned(len, alignment); - // SAFETY: We fully initialize this slice before freezing it. - unsafe { buffer.set_len(len) }; - Ok(WritableHostBuffer::new(Box::new( - DefaultWritableHostBuffer { buffer, alignment }, - ))) - } -} - -#[derive(Debug)] -struct DefaultWritableHostBuffer { - buffer: ByteBufferMut, - alignment: Alignment, -} - -#[derive(Debug)] -struct HostBufferOwner { - buffer: ByteBufferMut, -} - -impl AsRef<[u8]> for HostBufferOwner { - fn as_ref(&self) -> &[u8] { - self.buffer.as_slice() - } -} - -impl HostBufferMut for DefaultWritableHostBuffer { - fn len(&self) -> usize { - self.buffer.len() - } - - fn alignment(&self) -> Alignment { - self.alignment - } - - fn as_mut_slice(&mut self) -> &mut [u8] { - self.buffer.as_mut_slice() - } - - fn freeze(self: Box) -> ByteBuffer { - let Self { buffer, alignment } = *self; - let bytes = Bytes::from_owner(HostBufferOwner { buffer }); - ByteBuffer::from_bytes_aligned(bytes, alignment) - } -} - #[cfg(test)] mod tests { - use std::sync::Arc; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - - use super::*; + use vortex_buffer::BufferAllocatorRef; - #[derive(Debug)] - struct CountingAllocator { - allocations: Arc, - } - - impl HostAllocator for CountingAllocator { - fn allocate(&self, len: usize, alignment: Alignment) -> VortexResult { - self.allocations.fetch_add(1, Ordering::Relaxed); - DefaultHostAllocator.allocate(len, alignment) - } - } - - #[test] - fn writable_host_buffer_freeze_round_trip() { - let allocator = DefaultHostAllocator; - let mut writable = allocator.allocate(16, Alignment::new(8)).unwrap(); - for (idx, byte) in writable.as_mut_slice().iter_mut().enumerate() { - *byte = u8::try_from(idx).unwrap(); - } - - let host = writable.freeze(); - assert_eq!(host.len(), 16); - assert!(host.is_aligned(Alignment::new(8))); - assert_eq!(host.as_slice(), (0u8..16).collect::>().as_slice()); - } + use super::MemorySession; #[test] fn memory_session_replaces_allocator() { - let allocations = Arc::new(AtomicUsize::new(0)); - let allocator = Arc::new(CountingAllocator { - allocations: Arc::clone(&allocations), - }); + let allocator = BufferAllocatorRef::statically_allocated(); let mut session = MemorySession::default(); session.set_allocator(allocator); - drop(session.allocator().allocate(4, Alignment::none()).unwrap()); - assert_eq!(allocations.load(Ordering::Relaxed), 1); - } - - #[test] - fn typed_allocation_uses_type_alignment() { - let allocator = DefaultHostAllocator; - let writable = allocator.allocate_typed::(4).unwrap(); - assert_eq!(writable.len(), 4 * size_of::()); - assert_eq!(writable.alignment(), Alignment::of::()); - } - - #[test] - fn typed_mut_slice_round_trip() { - let allocator = DefaultHostAllocator; - let mut writable = allocator.allocate_typed::(4).unwrap(); - writable - .as_mut_slice_typed::() - .unwrap() - .copy_from_slice(&[10, 20, 30, 40]); - - let frozen = writable.freeze(); - let values = unsafe { - std::slice::from_raw_parts( - frozen.as_slice().as_ptr().cast::(), - frozen.len() / size_of::(), - ) - }; - assert_eq!(values, [10, 20, 30, 40]); - } - - #[test] - fn typed_mut_slice_rejects_length_mismatch() { - let allocator = DefaultHostAllocator; - let mut writable = allocator.allocate(7, Alignment::none()).unwrap(); - assert!(writable.as_mut_slice_typed::().is_err()); - } - - #[test] - fn freeze_typed_round_trip() { - let allocator = DefaultHostAllocator; - let mut writable = allocator.allocate_typed::(4).unwrap(); - writable - .as_mut_slice_typed::() - .unwrap() - .copy_from_slice(&[1, 3, 5, 7]); - - let frozen = writable.freeze_typed::().unwrap(); - assert_eq!(frozen.as_slice(), [1, 3, 5, 7]); - } - - #[test] - fn freeze_typed_rejects_length_mismatch() { - let allocator = DefaultHostAllocator; - let writable = allocator.allocate(7, Alignment::none()).unwrap(); - let err = writable.freeze_typed::().unwrap_err(); - let msg = format!("{err}"); - assert!(msg.contains("not a multiple of")); + let buffer = session.allocator().copy_from([1u32, 2, 3]); + assert_eq!(buffer.as_slice(), [1, 2, 3]); } } diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index 689cfab18a0..4285b79388c 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -144,7 +144,8 @@ impl ScalarFnVTable for VariantGet { .map_or(DType::Variant(Nullability::Nullable), DType::as_nullable); if !dtype.is_variant() { - let mut builder = builder_with_capacity_in(ctx.allocator(), &dtype, input.len()); + let mut builder = + builder_with_capacity_in(ctx.allocator().clone(), &dtype, input.len()); for idx in 0..input.len() { let scalar = input.execute_scalar(idx, ctx)?; let output = variant_get_scalar(&scalar, options, &dtype)?; diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 415ba4fc4b4..8b8dcc4e7df 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -45,6 +45,15 @@ impl BufferAllocatorRef { self.0.is_none() } + /// Returns true if both references point to the same allocator. + pub fn ptr_eq(&self, other: &Self) -> bool { + match (&self.0, &other.0) { + (None, None) => true, + (Some(lhs), Some(rhs)) => Arc::ptr_eq(lhs, rhs), + _ => false, + } + } + /// Create a mutable buffer with this allocator. pub fn with_capacity(&self, capacity: usize) -> BufferMut { BufferMut::with_capacity_in(capacity, self.clone()) diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index bddaff61a02..626a3fabffe 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -61,6 +61,7 @@ vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } [dev-dependencies] +allocator-api2 = { workspace = true } divan = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index dc009b973d0..b54da20e192 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -500,17 +500,20 @@ impl VortexOpenOptions { #[cfg(test)] mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; 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 futures::future::BoxFuture; use parking_lot::Mutex; use vortex_array::IntoArray; use vortex_array::buffer::BufferHandle; - use vortex_array::memory::DefaultHostAllocator; - use vortex_array::memory::HostAllocator; + use vortex_array::memory::BufferAllocatorRef; use vortex_array::memory::MemorySessionExt; - use vortex_array::memory::WritableHostBuffer; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; @@ -575,10 +578,16 @@ mod tests { allocations: Arc, } - impl HostAllocator for CountingAllocator { - fn allocate(&self, len: usize, alignment: Alignment) -> VortexResult { + // SAFETY: this forwards memory operations to Global and only counts allocations. + unsafe impl Allocator for CountingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { self.allocations.fetch_add(1, Ordering::Relaxed); - DefaultHostAllocator.allocate(len, alignment) + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: ptr and layout came from Global. + unsafe { Global.deallocate(ptr, layout) } } } @@ -783,7 +792,7 @@ mod tests { std::fs::write(&file_path, ByteBuffer::from(buf).as_slice()).unwrap(); let allocations = Arc::new(AtomicUsize::new(0)); - let session = session.with_allocator(Arc::new(CountingAllocator { + let session = session.with_allocator(BufferAllocatorRef::new(CountingAllocator { allocations: Arc::clone(&allocations), })); diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index 462bc498f82..ed36c2a02a1 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -17,8 +17,7 @@ use object_store::ObjectStore; use object_store::ObjectStoreExt; use object_store::path::Path as ObjectPath; use vortex_array::buffer::BufferHandle; -use vortex_array::memory::DefaultHostAllocator; -use vortex_array::memory::HostAllocatorRef; +use vortex_array::memory::BufferAllocatorRef; use vortex_buffer::Alignment; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -41,7 +40,7 @@ pub struct ObjectStoreReadAt { path: ObjectPath, uri: Arc, handle: Handle, - allocator: HostAllocatorRef, + allocator: BufferAllocatorRef, concurrency: usize, coalesce_config: Option, } @@ -49,7 +48,12 @@ pub struct ObjectStoreReadAt { impl ObjectStoreReadAt { /// Create a new object store source. pub fn new(store: Arc, path: ObjectPath, handle: Handle) -> Self { - Self::new_with_allocator(store, path, handle, Arc::new(DefaultHostAllocator)) + Self::new_with_allocator( + store, + path, + handle, + BufferAllocatorRef::statically_allocated(), + ) } /// Create a new object store source with a custom writable buffer allocator. @@ -57,7 +61,7 @@ impl ObjectStoreReadAt { store: Arc, path: ObjectPath, handle: Handle, - allocator: HostAllocatorRef, + allocator: BufferAllocatorRef, ) -> Self { let uri = Arc::from(path.to_string()); Self { @@ -88,7 +92,7 @@ async fn read_object_store_range( store: Arc, path: ObjectPath, io_handle: Handle, - allocator: HostAllocatorRef, + allocator: BufferAllocatorRef, request: ReadAtRequest, ) -> VortexResult { let ReadAtRequest { @@ -97,7 +101,9 @@ async fn read_object_store_range( alignment, } = request; let range = offset..(offset + length as u64); - let mut buffer = allocator.allocate(length, alignment)?; + let mut buffer = allocator.with_capacity_aligned::(length, alignment); + // SAFETY: each return path checks that every byte was initialized. + unsafe { buffer.set_len(length) }; let response = store .get_opts( @@ -188,7 +194,7 @@ impl VortexReadAt for ObjectStoreReadAt { let store = Arc::clone(&self.store); let path = self.path.clone(); let handle = self.handle.clone(); - let allocator = Arc::clone(&self.allocator); + let allocator = self.allocator.clone(); let io_handle = handle.clone(); handle .spawn_io(read_object_store_range( @@ -209,7 +215,7 @@ impl VortexReadAt for ObjectStoreReadAt { let store = Arc::clone(&self.store); let path = self.path.clone(); let handle = self.handle.clone(); - let allocator = Arc::clone(&self.allocator); + let allocator = self.allocator.clone(); let concurrency = self.concurrency.max(1); let (mut send, recv) = mpsc::channel(concurrency); let io_handle = handle.clone(); @@ -222,7 +228,7 @@ impl VortexReadAt for ObjectStoreReadAt { let store = Arc::clone(&store); let path = path.clone(); let io_handle = io_handle.clone(); - let allocator = Arc::clone(&allocator); + let allocator = allocator.clone(); async move { let result = read_object_store_range(store, path, io_handle, allocator, request).await; diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 3d59a595f70..21aac88923f 100644 --- a/vortex-io/src/std_file/read_at.rs +++ b/vortex-io/src/std_file/read_at.rs @@ -17,8 +17,7 @@ use std::sync::Arc; use futures::FutureExt; use futures::future::BoxFuture; use vortex_array::buffer::BufferHandle; -use vortex_array::memory::DefaultHostAllocator; -use vortex_array::memory::HostAllocatorRef; +use vortex_array::memory::BufferAllocatorRef; use vortex_buffer::Alignment; use vortex_error::VortexResult; @@ -66,20 +65,20 @@ pub struct FileReadAt { uri: Arc, file: Arc, handle: Handle, - allocator: HostAllocatorRef, + allocator: BufferAllocatorRef, } impl FileReadAt { /// Open a file for reading. pub fn open(path: impl AsRef, handle: Handle) -> VortexResult { - Self::open_with_allocator(path, handle, Arc::new(DefaultHostAllocator)) + Self::open_with_allocator(path, handle, BufferAllocatorRef::statically_allocated()) } /// Open a file for reading using a custom writable buffer allocator. pub fn open_with_allocator( path: impl AsRef, handle: Handle, - allocator: HostAllocatorRef, + allocator: BufferAllocatorRef, ) -> VortexResult { let path = path.as_ref(); let uri = path.to_string_lossy().to_string().into(); @@ -123,11 +122,13 @@ impl VortexReadAt for FileReadAt { ) -> BoxFuture<'static, VortexResult> { let file = Arc::clone(&self.file); let handle = self.handle.clone(); - let allocator = Arc::clone(&self.allocator); + let allocator = self.allocator.clone(); async move { handle .spawn_blocking(move || { - let mut buffer = allocator.allocate(length, alignment)?; + let mut buffer = allocator.with_capacity_aligned::(length, alignment); + // SAFETY: read_exact_at initializes every byte before the buffer is frozen. + unsafe { buffer.set_len(length) }; read_exact_at(&file, buffer.as_mut_slice(), offset)?; Ok(BufferHandle::new_host(buffer.freeze())) }) From eb02767b696a47202e331b6fd8953991d7fae097 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 27 Aug 2026 17:35:50 -0400 Subject: [PATCH 2/3] fix(array): keep execution context sync Signed-off-by: Nicholas Gates --- vortex-array/src/executor.rs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 04c4a31ce4a..dab9ccfcfb5 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -12,12 +12,12 @@ //! See for the full execution //! narrative, diagrams, and walkthroughs. -use std::cell::OnceCell; use std::env::VarError; use std::fmt; use std::fmt::Display; use std::sync::Arc; use std::sync::LazyLock; +use std::sync::OnceLock; #[cfg(debug_assertions)] use std::sync::atomic::AtomicUsize; #[cfg(debug_assertions)] @@ -351,8 +351,8 @@ struct StackFrame { #[derive(Debug, Clone)] pub struct ExecutionCtx { session: VortexSession, - // OnceCell avoids cloning the session allocator when a context does not allocate. - allocator: OnceCell, + // OnceLock avoids cloning the session allocator when a context does not allocate. + allocator: OnceLock, execute_parent_kernels: Arc, #[cfg(debug_assertions)] id: usize, @@ -370,7 +370,7 @@ impl ExecutionCtx { let execute_parent_kernels = session.kernels().execute_parent_snapshot(); Self { session, - allocator: OnceCell::new(), + allocator: OnceLock::new(), execute_parent_kernels, #[cfg(debug_assertions)] id: { @@ -394,7 +394,7 @@ impl ExecutionCtx { /// Set the allocator for this execution context. pub fn with_allocator(mut self, allocator: BufferAllocatorRef) -> Self { - self.allocator = OnceCell::from(allocator); + self.allocator = OnceLock::from(allocator); self } @@ -937,6 +937,7 @@ impl VortexSessionExecute for VortexSession { #[cfg(test)] mod tests { + use static_assertions::assert_impl_all; use vortex_session::SessionExt; use vortex_session::VortexSession; @@ -952,6 +953,8 @@ mod tests { use crate::optimizer::kernels::KernelSession; use crate::optimizer::kernels::execute_parent_key; + assert_impl_all!(ExecutionCtx: Send, Sync); + fn noop_execute_parent( _child: &ArrayRef, _parent: &ArrayRef, From 2e9ce475c196afc77c8225d6824fdf7537ca0741 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 16:50:28 -0400 Subject: [PATCH 3/3] fix(buffer): compare static allocator references Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 8b8dcc4e7df..fa6e4fea324 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -415,6 +415,17 @@ mod tests { } } + #[test] + fn allocator_identity() { + let static_allocator = BufferAllocatorRef::statically_allocated(); + assert!(static_allocator.ptr_eq(&BufferAllocatorRef::statically_allocated())); + + let custom_allocator = BufferAllocatorRef::new(TrackingAllocator::default()); + assert!(custom_allocator.ptr_eq(&custom_allocator.clone())); + assert!(!custom_allocator.ptr_eq(&static_allocator)); + assert!(!custom_allocator.ptr_eq(&BufferAllocatorRef::new(TrackingAllocator::default()))); + } + #[test] fn allocation_lives_until_last_view() { let allocator = TrackingAllocator::default();