From 7986ffe411adbddb6a2ea6828626b8a3c89819d1 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 27 Aug 2026 12:34:32 -0400 Subject: [PATCH 01/29] feat(buffer): add allocator-backed storage Signed-off-by: Nicholas Gates --- Cargo.lock | 1 + Cargo.toml | 1 + vortex-buffer/Cargo.toml | 1 + vortex-buffer/src/allocation.rs | 368 +++++++++++++++++++++++++++++ vortex-buffer/src/arrow.rs | 16 +- vortex-buffer/src/buffer.rs | 325 +++++++++++++++++-------- vortex-buffer/src/buffer_mut.rs | 407 +++++++++++++++++++++----------- vortex-buffer/src/lib.rs | 9 +- vortex-buffer/src/serde.rs | 2 +- 9 files changed, 883 insertions(+), 247 deletions(-) create mode 100644 vortex-buffer/src/allocation.rs diff --git a/Cargo.lock b/Cargo.lock index 3e8ed089b25..e35e6b84f00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10736,6 +10736,7 @@ dependencies = [ name = "vortex-buffer" version = "0.1.0" dependencies = [ + "allocator-api2", "arrow-buffer 59.2.0", "bitvec", "bytes", diff --git a/Cargo.toml b/Cargo.toml index f2e5ca7e000..842f4bc38b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -117,6 +117,7 @@ async-fs = "2.2.0" async-lock = "3.4" async-stream = "0.3.6" async-trait = "0.1.89" +allocator-api2 = "0.2.21" base16ct = "1.0.0" bigdecimal = "0.4.8" bindgen = "0.72.0" diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index 705c992f87d..c77d07f253d 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -23,6 +23,7 @@ serde = ["dep:serde", "serde/serde_derive"] warn-copy = ["dep:tracing"] [dependencies] +allocator-api2 = { workspace = true } arrow-buffer = { workspace = true } bitvec = { workspace = true } bytes = { workspace = true } diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs new file mode 100644 index 00000000000..54e62e0d58e --- /dev/null +++ b/vortex-buffer/src/allocation.rs @@ -0,0 +1,368 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Allocator-backed storage for Vortex buffers. + +use std::alloc::Layout; +use std::fmt; +use std::fmt::Debug; +use std::ptr::NonNull; +use std::sync::Arc; +use std::sync::LazyLock; + +use allocator_api2::alloc::AllocError; +use allocator_api2::alloc::Allocator; +use allocator_api2::alloc::Global; +use allocator_api2::alloc::handle_alloc_error; + +use crate::Alignment; +use crate::BufferMut; + +/// An allocator that can back a Vortex buffer. +pub trait BufferAllocator: Allocator + Debug + Send + Sync + 'static {} + +impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static {} + +/// A shared reference to a buffer allocator. +#[derive(Clone)] +pub struct BufferAllocatorRef(Arc); + +impl BufferAllocatorRef { + /// Wrap an allocator in a shared reference. + pub fn new(allocator: impl BufferAllocator) -> Self { + Self(Arc::new(allocator)) + } + + /// Return a shared reference to the static allocator. + pub fn statically_allocated() -> Self { + STATIC_ALLOCATOR.clone() + } + + /// Create a mutable buffer with this allocator. + pub fn with_capacity(&self, capacity: usize) -> BufferMut { + BufferMut::with_capacity_in(capacity, self.clone()) + } + + /// Create an aligned mutable buffer with this allocator. + pub fn with_capacity_aligned(&self, capacity: usize, alignment: Alignment) -> BufferMut { + BufferMut::with_capacity_aligned_in(capacity, alignment, self.clone()) + } + + /// Create a zeroed mutable buffer with this allocator. + pub fn zeroed(&self, len: usize) -> BufferMut { + BufferMut::zeroed_in(len, self.clone()) + } + + /// Copy values into a mutable buffer made by this allocator. + pub fn copy_from(&self, values: impl AsRef<[T]>) -> BufferMut { + BufferMut::copy_from_in(values, self.clone()) + } +} + +impl Debug for BufferAllocatorRef { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +// SAFETY: all calls are forwarded to the same allocator value held by the Arc. +unsafe impl Allocator for BufferAllocatorRef { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.0.allocate(layout) + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + self.0.allocate_zeroed(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow(ptr, old_layout, new_layout) } + } + + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } + } + + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.shrink(ptr, old_layout, new_layout) } + } +} + +/// The allocator used by buffer APIs that do not take an allocator. +#[derive(Clone, Copy, Debug, Default)] +pub struct StaticBufferAllocator; + +impl StaticBufferAllocator { + /// Create a mutable buffer with the static allocator. + pub fn with_capacity(capacity: usize) -> BufferMut { + BufferMut::with_capacity(capacity) + } + + /// Create an aligned mutable buffer with the static allocator. + pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> BufferMut { + BufferMut::with_capacity_aligned(capacity, alignment) + } + + /// Create a zeroed mutable buffer with the static allocator. + pub fn zeroed(len: usize) -> BufferMut { + BufferMut::zeroed(len) + } + + /// Copy values into a mutable buffer made by the static allocator. + pub fn copy_from(values: impl AsRef<[T]>) -> BufferMut { + BufferMut::copy_from(values) + } +} + +// SAFETY: Global satisfies the Allocator contract and this type only forwards to it. +unsafe impl Allocator for StaticBufferAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + Global.allocate(layout) + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + Global.allocate_zeroed(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.deallocate(ptr, layout) } + } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + + unsafe fn grow_zeroed( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } + } + + unsafe fn shrink( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.shrink(ptr, old_layout, new_layout) } + } +} + +static STATIC_ALLOCATOR: LazyLock = + LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); + +pub(crate) struct Allocation { + ptr: NonNull, + layout: Layout, + capacity: usize, + allocator: BufferAllocatorRef, +} + +// SAFETY: Allocation owns its memory, and its allocator is Send + Sync. +unsafe impl Send for Allocation {} +// SAFETY: shared access to Allocation never permits mutation of the allocation. +unsafe impl Sync for Allocation {} + +impl Allocation { + pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, false) + } + + pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, true) + } + + fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { + if layout.size() == 0 { + return Self { + ptr: layout.dangling_ptr(), + layout, + capacity: 0, + allocator, + }; + } + + let allocation = if zeroed { + allocator.allocate_zeroed(layout) + } else { + allocator.allocate(layout) + } + .unwrap_or_else(|_| handle_alloc_error(layout)); + + Self { + ptr: allocation.cast(), + layout, + capacity: allocation.len(), + allocator, + } + } + + pub(crate) fn ptr(&self) -> NonNull { + self.ptr + } + + pub(crate) fn size(&self) -> usize { + self.capacity + } + + pub(crate) fn alignment(&self) -> Alignment { + Alignment::new(self.layout.align()) + } + + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + &self.allocator + } + + pub(crate) fn grow(&mut self, layout: Layout) { + if self.layout.size() == 0 { + *self = Self::allocate(layout, self.allocator.clone()); + return; + } + + let allocation = + // SAFETY: ptr and layout describe a live block allocated by self.allocator. The new + // layout is at least as large as the old layout. + unsafe { self.allocator.grow(self.ptr, self.layout, layout) } + .unwrap_or_else(|_| handle_alloc_error(layout)); + self.ptr = allocation.cast(); + self.layout = layout; + self.capacity = allocation.len(); + } +} + +impl Drop for Allocation { + fn drop(&mut self) { + if self.layout.size() == 0 { + return; + } + // SAFETY: ptr and layout describe a live block allocated by self.allocator. + unsafe { self.allocator.deallocate(self.ptr, self.layout) } + } +} + +pub(crate) trait BufferOwner: Send + Sync + 'static { + fn as_slice(&self) -> &[u8]; +} + +impl BufferOwner for T +where + T: AsRef<[u8]> + Send + Sync + 'static, +{ + fn as_slice(&self) -> &[u8] { + self.as_ref() + } +} + +pub(crate) enum BufferBacking { + Owned(Allocation), + External { _owner: Box }, +} + +impl BufferBacking { + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + match self { + Self::Owned(allocation) => allocation.allocator(), + Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), + } + } +} + +#[cfg(test)] +mod tests { + use std::alloc::Layout; + use std::ptr::NonNull; + use std::sync::Arc; + 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 crate::Alignment; + use crate::BufferAllocatorRef; + + #[derive(Clone, Debug, Default)] + struct TrackingAllocator { + state: Arc, + } + + #[derive(Debug, Default)] + struct TrackingState { + allocations: AtomicUsize, + deallocations: AtomicUsize, + alignment: AtomicUsize, + } + + // SAFETY: this forwards all memory operations to Global and only records call metadata. + unsafe impl Allocator for TrackingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.state.allocations.fetch_add(1, Ordering::Relaxed); + self.state + .alignment + .store(layout.align(), Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + self.state.deallocations.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller passes the pointer and layout returned by Global. + unsafe { Global.deallocate(ptr, layout) } + } + } + + #[test] + fn allocation_lives_until_last_view() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let buffer = BufferAllocatorRef::new(allocator) + .copy_from([1u32, 2, 3, 4]) + .freeze(); + let view = buffer.slice(0..2); + + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!( + state.alignment.load(Ordering::Relaxed), + *Alignment::DEFAULT_ALIGNMENT + ); + drop(buffer); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + drop(view); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + } +} diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index aa96cab2e84..f63aaadaff0 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -35,12 +35,8 @@ impl Buffer { ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + debug_assert_eq!(length, bytes.len() / size_of::()); + Self::from_bytes_aligned(bytes, alignment) } /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`. @@ -74,12 +70,8 @@ impl ByteBuffer { ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + debug_assert_eq!(length, bytes.len()); + Self::from_bytes_aligned(bytes, alignment) } } diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 8c78b9f8a1c..4b5ee8ef6fd 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -8,9 +8,11 @@ use std::fmt::Debug; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; -use std::marker::PhantomData; use std::ops::Deref; use std::ops::RangeBounds; +use std::ptr::NonNull; +use std::sync::Arc; +use std::sync::LazyLock; use bytes::Buf; use bytes::Bytes; @@ -18,6 +20,9 @@ use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; +use crate::BufferAllocatorRef; +use crate::BufferBacking; use crate::BufferMut; use crate::ByteBuffer; use crate::debug::TruncatedDebug; @@ -26,63 +31,109 @@ use crate::trusted_len::TrustedLen; /// An immutable buffer of items of `T`. #[derive(Clone)] pub struct Buffer { - pub(crate) bytes: Bytes, + pub(crate) ptr: NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, - pub(crate) _marker: PhantomData, + pub(crate) backing: Arc, } +// SAFETY: Buffer is an immutable view over backing memory. Its pointer remains valid while the +// backing is live, and sharing elements follows the same bounds as sharing a slice. +unsafe impl Send for Buffer {} +// SAFETY: see the Send implementation above. +unsafe impl Sync for Buffer {} + /// Zero-length backing for empty buffers, "aligned" to [`Alignment::MAX`] so it satisfies any /// valid alignment without allocating. A zero-length slice never reads memory, so it may use a /// dangling pointer as long as it is non-null and aligned. -const EMPTY_BACKING: &[u8] = { +const EMPTY_BYTES: &[u8] = { let ptr = std::ptr::without_provenance(Alignment::MAX.as_usize()); // SAFETY: the pointer is non-null and aligned, and the slice is zero-length. unsafe { std::slice::from_raw_parts(ptr, 0) } }; +static EMPTY_BACKING: LazyLock> = LazyLock::new(|| { + Arc::new(BufferBacking::External { + _owner: Box::new(EMPTY_BYTES), + }) +}); + impl Default for Buffer { fn default() -> Self { Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment: Alignment::of::(), - _marker: PhantomData, + backing: EMPTY_BACKING.clone(), } } } -impl PartialEq for Buffer { +impl PartialEq for Buffer { #[inline] fn eq(&self, other: &Self) -> bool { - self.bytes == other.bytes + self.as_slice() == other.as_slice() } } -impl Eq for Buffer {} +impl Eq for Buffer {} -impl Ord for Buffer { +impl Ord for Buffer { #[inline] fn cmp(&self, other: &Self) -> Ordering { - self.bytes.cmp(&other.bytes) + self.as_slice().cmp(other.as_slice()) } } -impl PartialOrd for Buffer { +impl PartialOrd for Buffer { #[inline] fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) + self.as_slice().partial_cmp(other.as_slice()) } } -impl Hash for Buffer { +impl Hash for Buffer { #[inline] fn hash(&self, state: &mut H) { - self.bytes.as_ref().hash(state) + self.as_slice().hash(state) } } impl Buffer { + pub(crate) fn from_allocation( + allocation: Allocation, + offset: usize, + length: usize, + alignment: Alignment, + ) -> Self { + // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; + Self { + ptr, + length, + alignment, + backing: Arc::new(BufferBacking::Owned(allocation)), + } + } + + fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self { + let bytes = owner.as_slice(); + let length = bytes.len() / size_of::(); + let ptr = if length == 0 { + empty_ptr() + } else { + NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") + }; + Self { + ptr, + length, + alignment, + backing: Arc::new(BufferBacking::External { + _owner: Box::new(owner), + }), + } + } + /// Returns a new `Buffer` copied from the provided `Vec`, `&[T]`, etc. /// /// Due to our underlying usage of `bytes::Bytes`, we are unable to take zero-copy ownership @@ -93,6 +144,11 @@ impl Buffer { BufferMut::copy_from(values).freeze() } + /// Returns a new `Buffer` copied with the provided allocator. + pub fn copy_from_in(values: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + BufferMut::copy_from_in(values, allocator).freeze() + } + /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -120,6 +176,11 @@ impl Buffer { Self::zeroed_aligned(len, Alignment::of::()) } + /// Create a new zeroed `Buffer` with the provided allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + BufferMut::zeroed_in(len, allocator).freeze() + } + /// Create a new zeroed `Buffer` with the requested alignment. /// /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than @@ -160,10 +221,10 @@ impl Buffer { ); } Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment, - _marker: PhantomData, + backing: EMPTY_BACKING.clone(), } } @@ -175,6 +236,14 @@ impl Buffer { BufferMut::full(item, len).freeze() } + /// Create a full `Buffer` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + BufferMut::full_in(item, len, allocator).freeze() + } + /// Create a `Buffer` zero-copy from a `ByteBuffer`. /// /// ## Panics @@ -193,7 +262,29 @@ impl Buffer { /// Panics if the buffer is not aligned to the given alignment, if the length is not a multiple /// of the size of `T`, or if the given alignment is not aligned to that of `T`. pub fn from_byte_buffer_aligned(buffer: ByteBuffer, alignment: Alignment) -> Self { - Self::from_bytes_aligned(buffer.into_inner(), alignment) + if !alignment.is_aligned_to(Alignment::of::()) { + vortex_panic!( + "Alignment {} must be compatible with the scalar type's alignment {}", + alignment, + Alignment::of::(), + ); + } + if !alignment.is_ptr_aligned(buffer.as_ptr()) { + vortex_panic!("Buffer must align to the requested alignment {}", alignment); + } + if !buffer.len().is_multiple_of(size_of::()) { + vortex_panic!( + "Buffer length {} must be a multiple of the scalar type's size {}", + buffer.len(), + size_of::() + ); + } + Self { + ptr: buffer.ptr.cast(), + length: buffer.length / size_of::(), + alignment, + backing: buffer.backing, + } } /// Create a `Buffer` zero-copy from a `Bytes`. @@ -223,13 +314,7 @@ impl Buffer { size_of::() ); } - let length = bytes.len() / size_of::(); - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + Self::from_owner(bytes, alignment) } /// Create a buffer with values from the TrustedLen iterator. @@ -248,7 +333,8 @@ impl Buffer { Ok(mut_buf) => mut_buf.map_each_in_place(f), Err(buf) => { let len = buf.len(); - let mut out_buf = BufferMut::with_capacity(len); + let allocator = buf.backing.allocator().clone(); + let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf .spare_capacity_mut() .iter_mut() @@ -265,7 +351,6 @@ impl Buffer { /// Clear the buffer, preserving existing capacity. pub fn clear(&mut self) { - self.bytes.clear(); self.length = 0; } @@ -290,19 +375,35 @@ impl Buffer { self.alignment } + /// Returns the allocator to use for derived buffers. + /// + /// External buffers use the static allocator. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.backing.allocator() + } + + /// Returns a raw pointer to the buffer's data. + #[inline(always)] + pub fn as_ptr(&self) -> *const T { + self.ptr.as_ptr() + } + /// Returns a slice over the buffer of elements of type T. #[allow(clippy::inline_always)] #[inline(always)] pub fn as_slice(&self) -> &[T] { - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(self.bytes.as_ptr().cast(), self.length) } + // SAFETY: ptr points into the live backing and construction checks its alignment. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } } /// Return a view over the buffer as an opaque byte slice. #[allow(clippy::inline_always)] #[inline(always)] pub fn as_bytes(&self) -> &[u8] { - self.bytes.as_ref() + // SAFETY: the element range is initialized and remains live through backing. + unsafe { + std::slice::from_raw_parts(self.ptr.as_ptr().cast(), size_of_val(self.as_slice())) + } } /// Returns an iterator over the buffer of elements of type T. @@ -378,8 +479,6 @@ impl Buffer { } let begin_byte = begin * size_of::(); - let end_byte = end * size_of::(); - if !alignment.is_offset_aligned(begin_byte) { vortex_panic!( "range start must be aligned to {alignment:?}, byte {}", @@ -391,10 +490,11 @@ impl Buffer { } Self { - bytes: self.bytes.slice(begin_byte..end_byte), + // SAFETY: begin is in bounds and the alignment check applies to the new pointer. + ptr: unsafe { self.ptr.add(begin) }, length: end - begin, alignment, - _marker: Default::default(), + backing: Arc::clone(&self.backing), } } @@ -434,74 +534,90 @@ impl Buffer { vortex_panic!("slice_ref subset must be aligned to {:?}", alignment); } - let subset_u8 = - unsafe { std::slice::from_raw_parts(subset.as_ptr().cast(), size_of_val(subset)) }; + let start = self.as_ptr().addr(); + let end = start + size_of_val(self.as_slice()); + let subset_start = subset.as_ptr().addr(); + let subset_end = subset_start + .checked_add(size_of_val(subset)) + .vortex_expect("slice_ref address overflow"); + if subset_start < start || subset_end > end { + vortex_panic!("slice_ref subset must be contained in the buffer"); + } Self { - bytes: self.bytes.slice_ref(subset_u8), + ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"), length: subset.len(), alignment, - _marker: Default::default(), + backing: Arc::clone(&self.backing), } } - /// Returns the underlying aligned buffer. - pub fn inner(&self) -> &Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - &self.bytes - } - - /// Returns the underlying aligned buffer. + /// Returns the underlying bytes without copying. pub fn into_inner(self) -> Bytes { - debug_assert_eq!( - self.length * size_of::(), - self.bytes.len(), - "Own length has to be the same as the underlying bytes length" - ); - self.bytes + Bytes::from_owner(BufferBytesOwner { + ptr: self.ptr.cast(), + length: self.length * size_of::(), + backing: self.backing, + }) } /// Return the ByteBuffer for this `Buffer`. pub fn into_byte_buffer(self) -> ByteBuffer { ByteBuffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, - _marker: Default::default(), + backing: self.backing, } } /// Try to convert self into `BufferMut` if there is only a single strong reference. pub fn try_into_mut(self) -> Result, Self> { - self.bytes - .try_into_mut() - .map(|bytes| BufferMut { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) - .map_err(|bytes| Self { - bytes, - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - }) + let Self { + ptr, + length, + alignment, + backing, + } = self; + match Arc::try_unwrap(backing) { + Ok(BufferBacking::Owned(allocation)) => { + let offset = ptr.addr().get() - allocation.ptr().addr().get(); + let capacity = (allocation.size() - offset) / size_of::(); + Ok(BufferMut { + allocation, + offset, + length, + capacity, + alignment, + _marker: Default::default(), + }) + } + Ok(backing) => Err(Self { + ptr, + length, + alignment, + backing: Arc::new(backing), + }), + Err(backing) => Err(Self { + ptr, + length, + alignment, + backing, + }), + } } /// Convert self into `BufferMut`, cloning the data if there are multiple strong references. pub fn into_mut(self) -> BufferMut { - self.try_into_mut() - .unwrap_or_else(|buffer| BufferMut::::copy_from_aligned(&buffer, buffer.alignment)) + self.try_into_mut().unwrap_or_else(|buffer| { + let allocator = buffer.backing.allocator().clone(); + BufferMut::::copy_from_aligned_in(&buffer, buffer.alignment, allocator) + }) } /// Returns whether a `Buffer` is aligned to the given alignment. pub fn is_aligned(&self, alignment: Alignment) -> bool { - alignment.is_ptr_aligned(self.bytes.as_ptr()) + alignment.is_ptr_aligned(self.as_ptr()) } /// Return a `Buffer` with the given alignment. Where possible, this will be zero-copy. @@ -517,7 +633,8 @@ impl Buffer { "Buffer is not aligned to requested alignment {alignment}, copying: {bt}" ) } - Self::copy_from_aligned(self, alignment) + let allocator = self.backing.allocator().clone(); + BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze() } } @@ -553,10 +670,10 @@ impl Buffer { ); Buffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, - _marker: PhantomData, + backing: self.backing, } } } @@ -651,34 +768,21 @@ impl AsRef<[u8]> for Wrapper { impl From> for Buffer where - T: Send + 'static, + T: Send + Sync + 'static, { fn from(value: Vec) -> Self { - let original_len = value.len(); let wrapped_vec = Wrapper(value); - - let bytes = Bytes::from_owner(wrapped_vec); - - assert_eq!(bytes.as_ptr().align_offset(align_of::()), 0); - - Self { - bytes, - length: original_len, - alignment: Alignment::of::(), - _marker: PhantomData, - } + assert_eq!( + wrapped_vec.as_ref().as_ptr().align_offset(align_of::()), + 0 + ); + Self::from_owner(wrapped_vec, Alignment::of::()) } } impl From for ByteBuffer { fn from(bytes: Bytes) -> Self { - let length = bytes.len(); - Self { - bytes, - length, - alignment: Alignment::of::(), - _marker: Default::default(), - } + Self::from_owner(bytes, Alignment::of::()) } } @@ -702,11 +806,36 @@ impl Buf for ByteBuffer { self.alignment ); } - self.bytes.advance(cnt); + assert!(cnt <= self.length, "cannot advance past the buffer length"); + // SAFETY: cnt is within the initialized byte range. + self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; } } +struct BufferBytesOwner { + ptr: NonNull, + length: usize, + backing: Arc, +} + +// SAFETY: the owner exposes immutable initialized bytes and keeps their backing live. +unsafe impl Send for BufferBytesOwner {} +unsafe impl Sync for BufferBytesOwner {} + +impl AsRef<[u8]> for BufferBytesOwner { + fn as_ref(&self) -> &[u8] { + let _ = &self.backing; + // SAFETY: ptr and length came from a live Buffer. + unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.length) } + } +} + +fn empty_ptr() -> NonNull { + let addr = 1usize << (usize::BITS - 1); + NonNull::new(std::ptr::without_provenance_mut(addr)).vortex_expect("empty pointer is non-null") +} + /// Owned iterator over a [`Buffer`]. pub struct BufferIterator { // Keep the buffer alive for the duration of the iteration. diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 4cd51bf2546..84699ab9695 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use core::mem::MaybeUninit; +use std::alloc::Layout; use std::any::type_name; use std::cmp::max; use std::fmt::Debug; @@ -12,23 +13,25 @@ use std::ops::DerefMut; use bytes::Buf; use bytes::BufMut; -use bytes::BytesMut; use bytes::buf::UninitSlice; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use crate::Alignment; +use crate::Allocation; use crate::Buffer; +use crate::BufferAllocatorRef; use crate::ByteBufferMut; use crate::debug::TruncatedDebug; use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. -#[derive(PartialEq, Eq)] pub struct BufferMut { - pub(crate) bytes: BytesMut, + pub(crate) allocation: Allocation, + pub(crate) offset: usize, pub(crate) length: usize, + pub(crate) capacity: usize, pub(crate) alignment: Alignment, pub(crate) _marker: std::marker::PhantomData, } @@ -36,7 +39,12 @@ pub struct BufferMut { impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. pub fn with_capacity(capacity: usize) -> Self { - Self::with_capacity_aligned(capacity, Alignment::of::()) + Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new `BufferMut` with the requested capacity and allocator. + pub fn with_capacity_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(capacity, Alignment::of::(), allocator) } /// Create a new `BufferMut` with the requested alignment and capacity. @@ -46,10 +54,24 @@ impl BufferMut { /// /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { - Self::with_capacity_preferred_aligned( + Self::with_capacity_aligned_in( + capacity, + alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new `BufferMut` with the requested alignment, capacity, and allocator. + pub fn with_capacity_aligned_in( + capacity: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::with_capacity_preferred_aligned_in( capacity, alignment, Some(Alignment::DEFAULT_ALIGNMENT), + allocator, ) } @@ -61,6 +83,21 @@ impl BufferMut { capacity: usize, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::with_capacity_preferred_aligned_in( + capacity, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a new allocator-backed `BufferMut` with a requested and preferred alignment. + pub fn with_capacity_preferred_aligned_in( + capacity: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { let actual = max( alignment, @@ -75,12 +112,19 @@ impl BufferMut { ); } - let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + actual.as_usize()); - bytes.align_empty(actual); + let size = capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let layout = Layout::from_size_align(size, actual.as_usize()) + .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); + let allocation = Allocation::allocate(layout, allocator); + let capacity = allocation.size() / size_of::(); Self { - bytes, + allocation, + offset: 0, length: 0, + capacity, alignment, _marker: Default::default(), } @@ -88,7 +132,12 @@ impl BufferMut { /// Create a new zeroed `BufferMut`. pub fn zeroed(len: usize) -> Self { - Self::zeroed_aligned(len, Alignment::of::()) + Self::zeroed_in(len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a new zeroed `BufferMut` with the requested allocator. + pub fn zeroed_in(len: usize, allocator: BufferAllocatorRef) -> Self { + Self::zeroed_aligned_in(len, Alignment::of::(), allocator) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -98,7 +147,21 @@ impl BufferMut { /// /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::zeroed_aligned_in(len, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create a zeroed `BufferMut` with an alignment and allocator. + pub fn zeroed_aligned_in( + len: usize, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a new zeroed `BufferMut` with the requested alignment. @@ -109,16 +172,36 @@ impl BufferMut { len: usize, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::zeroed_preferred_aligned_in( + len, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Create a zeroed allocator-backed buffer with a requested and preferred alignment. + pub fn zeroed_preferred_aligned_in( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { let preferred_alignment = preferred_alignment.unwrap_or(Alignment::of::()); let actual_alignment = max(preferred_alignment, alignment); - let mut bytes = BytesMut::zeroed((len * size_of::()) + actual_alignment.as_usize()); - bytes.advance(bytes.as_ptr().align_offset(actual_alignment.as_usize())); - unsafe { bytes.set_len(len * size_of::()) }; - let actual_len = bytes.len().checked_div(size_of::()).unwrap_or(0); + let size = len + .checked_mul(size_of::()) + .vortex_expect("buffer length overflow"); + let layout = Layout::from_size_align(size, actual_alignment.as_usize()) + .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); + let allocation = Allocation::allocate_zeroed(layout, allocator); + let capacity = allocation.size() / size_of::(); Self { - bytes, - length: actual_len, + allocation, + offset: 0, + length: len, + capacity, alignment, _marker: Default::default(), } @@ -136,7 +219,12 @@ impl BufferMut { /// /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned pub fn empty_aligned(alignment: Alignment) -> Self { - Self::empty_preferred_aligned(alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::empty_aligned_in(alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Create an empty `BufferMut` with an alignment and allocator. + pub fn empty_aligned_in(alignment: Alignment, allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_aligned_in(0, alignment, allocator) } /// Create a new empty `BufferMut` with the provided alignment. @@ -147,7 +235,12 @@ impl BufferMut { alignment: Alignment, preferred_alignment: Option, ) -> Self { - BufferMut::with_capacity_preferred_aligned(0, alignment, preferred_alignment) + BufferMut::with_capacity_preferred_aligned_in( + 0, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) } /// Create a new full `BufferMut` with the given value. @@ -155,14 +248,27 @@ impl BufferMut { where T: Copy, { - let mut buffer = BufferMut::::with_capacity(len); + Self::full_in(item, len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a full `BufferMut` with the given value and allocator. + pub fn full_in(item: T, len: usize, allocator: BufferAllocatorRef) -> Self + where + T: Copy, + { + let mut buffer = BufferMut::::with_capacity_in(len, allocator); buffer.push_n(item, len); buffer } /// Create a mutable scalar buffer by copying the contents of the slice. pub fn copy_from(other: impl AsRef<[T]>) -> Self { - Self::copy_from_aligned(other, Alignment::of::()) + Self::copy_from_in(other, BufferAllocatorRef::statically_allocated()) + } + + /// Create a mutable scalar buffer by copying with the given allocator. + pub fn copy_from_in(other: impl AsRef<[T]>, allocator: BufferAllocatorRef) -> Self { + Self::copy_from_aligned_in(other, Alignment::of::(), allocator) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -176,7 +282,21 @@ impl BufferMut { /// /// Panics when the requested alignment isn't itself aligned to type T. pub fn copy_from_aligned(other: impl AsRef<[T]>, alignment: Alignment) -> Self { - Self::copy_from_preferred_aligned(other, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + Self::copy_from_aligned_in(other, alignment, BufferAllocatorRef::statically_allocated()) + } + + /// Copy values into a mutable buffer with the given alignment and allocator. + pub fn copy_from_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + allocator, + ) } /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. @@ -191,13 +311,32 @@ impl BufferMut { other: impl AsRef<[T]>, alignment: Alignment, preferred_alignment: Option, + ) -> Self { + Self::copy_from_preferred_aligned_in( + other, + alignment, + preferred_alignment, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Copy values with the given allocator, requested alignment, and preferred alignment. + pub fn copy_from_preferred_aligned_in( + other: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + allocator: BufferAllocatorRef, ) -> Self { if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!("Given alignment is not aligned to type T") } let other = other.as_ref(); - let mut buffer = - Self::with_capacity_preferred_aligned(other.len(), alignment, preferred_alignment); + let mut buffer = Self::with_capacity_preferred_aligned_in( + other.len(), + alignment, + preferred_alignment, + allocator, + ); buffer.extend_from_slice(other); debug_assert_eq!(buffer.alignment(), alignment); buffer @@ -210,11 +349,15 @@ impl BufferMut { self.alignment } + /// Returns the allocator that owns this buffer. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.allocation.allocator() + } + /// Returns the length of the buffer. #[allow(clippy::inline_always)] #[inline(always)] pub fn len(&self) -> usize { - debug_assert_eq!(self.length, self.bytes.len() / size_of::()); self.length } @@ -228,29 +371,38 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.bytes.capacity() / size_of::() + self.capacity + } + + /// Returns a raw pointer to the buffer's data. + pub fn as_ptr(&self) -> *const T { + // SAFETY: offset always remains within the allocation. + unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } + } + + /// Returns a mutable raw pointer to the buffer's data. + pub fn as_mut_ptr(&mut self) -> *mut T { + // SAFETY: BufferMut uniquely owns the allocation and offset is in bounds. + unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_slice(&self) -> &[T] { - let raw_slice = self.bytes.as_ref(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts(raw_slice.as_ptr().cast(), self.length) } + // SAFETY: the allocation is live, offset is in bounds, and construction checks alignment. + unsafe { std::slice::from_raw_parts(self.as_ptr(), self.length) } } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_mut_slice(&mut self) -> &mut [T] { - let raw_slice = self.bytes.as_mut(); - // SAFETY: alignment of Buffer is checked on construction - unsafe { std::slice::from_raw_parts_mut(raw_slice.as_mut_ptr().cast(), self.length) } + // SAFETY: BufferMut uniquely owns the allocation and the initialized range is in bounds. + unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.length) } } /// Clear the buffer, retaining any existing capacity. #[inline] pub fn clear(&mut self) { - unsafe { self.bytes.set_len(0) } self.length = 0; } @@ -272,8 +424,7 @@ impl BufferMut { /// Reserves capacity for at least `additional` more elements to be inserted in the buffer. #[inline] pub fn reserve(&mut self, additional: usize) { - let additional_bytes = additional * size_of::(); - if additional_bytes <= self.bytes.capacity() - self.bytes.len() { + if additional <= self.capacity - self.length { // We can fit the additional bytes in the remaining capacity. Nothing to do. return; } @@ -282,18 +433,37 @@ impl BufferMut { self.reserve_allocate(additional); } - /// A separate function so we can inline the reserve call's fast path. According to `BytesMut` - /// this has significant performance implications. + /// A separate function so we can inline the reserve call's fast path. fn reserve_allocate(&mut self, additional: usize) { - let new_capacity: usize = - ((self.length + additional) * size_of::()) + self.alignment.as_usize(); + let required = self + .length + .checked_add(additional) + .vortex_expect("buffer capacity overflow"); // Make sure we at least double in size each time we re-allocate to amortize the cost - let new_capacity = new_capacity.max(self.bytes.capacity() * 2); - - let mut bytes = BytesMut::with_capacity(new_capacity); - bytes.align_empty(self.alignment); - bytes.extend_from_slice(&self.bytes); - self.bytes = bytes; + let new_capacity = required.max(self.capacity.saturating_mul(2)); + let new_size = new_capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + let physical_alignment = max(self.alignment, self.allocation.alignment()); + let layout = Layout::from_size_align(new_size, physical_alignment.as_usize()) + .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); + + if self.offset == 0 { + self.allocation.grow(layout); + } else { + let mut allocation = Allocation::allocate(layout, self.allocation.allocator().clone()); + // SAFETY: both ranges are valid for the initialized byte length and do not overlap. + unsafe { + std::ptr::copy_nonoverlapping( + self.as_ptr().cast::(), + allocation.ptr().as_ptr(), + self.length * size_of::(), + ); + } + std::mem::swap(&mut self.allocation, &mut allocation); + self.offset = 0; + } + self.capacity = self.allocation.size() / size_of::(); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -333,13 +503,9 @@ impl BufferMut { /// ``` #[inline] pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit] { - let dst = self.bytes.spare_capacity_mut().as_mut_ptr(); - unsafe { - std::slice::from_raw_parts_mut( - dst as *mut MaybeUninit, - self.capacity() - self.length, - ) - } + // SAFETY: offset + length is within the allocation and points at spare capacity. + let dst = unsafe { self.as_mut_ptr().add(self.length) }.cast::>(); + unsafe { std::slice::from_raw_parts_mut(dst, self.capacity() - self.length) } } /// Sets the length of the buffer. @@ -353,7 +519,6 @@ impl BufferMut { #[inline] pub unsafe fn set_len(&mut self, len: usize) { debug_assert!(len <= self.capacity()); - unsafe { self.bytes.set_len(len * size_of::()) }; self.length = len; } @@ -373,9 +538,8 @@ impl BufferMut { pub unsafe fn push_unchecked(&mut self, item: T) { // SAFETY: the caller ensures we have sufficient capacity unsafe { - let dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let dst = self.as_mut_ptr().add(self.length); dst.write(item); - self.bytes.set_len(self.bytes.len() + size_of::()) } self.length += 1; } @@ -402,7 +566,8 @@ impl BufferMut { where T: Copy, { - let mut dst: *mut T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + // SAFETY: the caller guarantees enough spare capacity. + let mut dst = unsafe { self.as_mut_ptr().add(self.length) }; // SAFETY: we checked the capacity in the reserve call unsafe { let end = dst.add(n); @@ -410,7 +575,6 @@ impl BufferMut { dst.write(item); dst = dst.add(1); } - self.bytes.set_len(self.bytes.len() + (n * size_of::())); } self.length += n; } @@ -430,19 +594,21 @@ impl BufferMut { #[inline] pub fn extend_from_slice(&mut self, slice: &[T]) { self.reserve(slice.len()); - let raw_slice = - unsafe { std::slice::from_raw_parts(slice.as_ptr().cast(), size_of_val(slice)) }; - self.bytes.extend_from_slice(raw_slice); + // SAFETY: reserve made the destination valid and non-overlapping for slice.len() values. + unsafe { + std::ptr::copy_nonoverlapping( + slice.as_ptr(), + self.as_mut_ptr().add(self.length), + slice.len(), + ); + } self.length += slice.len(); } /// Splits the buffer into two at the given index. /// /// Afterward, self contains elements `[0, at)`, and the returned buffer contains elements - /// `[at, capacity)`. It’s guaranteed that the memory does not move, that is, the address of - /// self does not change, and the address of the returned slice is at bytes after that. - /// - /// This is an O(1) operation that just increases the reference count and sets a few indices. + /// `[at, capacity)`. The returned buffer uses a new allocation. /// /// Panics if either half would have a length that is not a multiple of the alignment. pub fn split_off(&mut self, at: usize) -> Self { @@ -459,27 +625,25 @@ impl BufferMut { ); } - let new_bytes = self.bytes.split_off(bytes_at); - - // Adjust the lengths, given that length may be < at let new_length = self.length.saturating_sub(at); + let new_capacity = self.capacity - at; + let mut other = Self::with_capacity_aligned_in( + new_capacity, + self.alignment, + self.allocation.allocator().clone(), + ); + if new_length > 0 { + other.extend_from_slice(&self.as_slice()[at..]); + } self.length = self.length.min(at); + self.capacity = at; - BufferMut { - bytes: new_bytes, - length: new_length, - alignment: self.alignment, - _marker: Default::default(), - } + other } /// Absorbs a mutable buffer that was previously split off. /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.extend_from_slice(other.as_ref()). + /// This appends the contents of `other` to this buffer. pub fn unsplit(&mut self, other: Self) { if self.alignment != other.alignment { vortex_panic!( @@ -488,15 +652,16 @@ impl BufferMut { other.alignment ); } - self.bytes.unsplit(other.bytes); - self.length += other.length; + self.extend_from_slice(other.as_slice()); } /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { ByteBufferMut { - bytes: self.bytes, + allocation: self.allocation, + offset: self.offset, length: self.length * size_of::(), + capacity: self.capacity * size_of::(), alignment: self.alignment, _marker: Default::default(), } @@ -504,12 +669,7 @@ impl BufferMut { /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { - Buffer { - bytes: self.bytes.freeze(), - length: self.length, - alignment: self.alignment, - _marker: Default::default(), - } + Buffer::from_allocation(self.allocation, self.offset, self.length, self.alignment) } /// Map each element of the buffer with a closure. @@ -537,14 +697,10 @@ impl BufferMut { /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { if self.as_ptr().align_offset(alignment.as_usize()) == 0 { - Self { - bytes: self.bytes, - length: self.length, - alignment, - _marker: std::marker::PhantomData, - } + Self { alignment, ..self } } else { - Self::copy_from_aligned(self, alignment) + let allocator = self.allocation.allocator().clone(); + Self::copy_from_aligned_in(self, alignment, allocator) } } @@ -568,8 +724,10 @@ impl BufferMut { ); BufferMut { - bytes: self.bytes, + allocation: self.allocation, + offset: self.offset, length: self.length, + capacity: self.capacity, alignment: self.alignment, _marker: std::marker::PhantomData, } @@ -578,14 +736,24 @@ impl BufferMut { impl Clone for BufferMut { fn clone(&self) -> Self { - // NOTE(ngates): we cannot derive Clone since BytesMut copies on clone and the alignment - // might be messed up. - let mut buffer = BufferMut::::with_capacity_aligned(self.capacity(), self.alignment); + let mut buffer = BufferMut::::with_capacity_aligned_in( + self.capacity(), + self.alignment, + self.allocation.allocator().clone(), + ); buffer.extend_from_slice(self.as_slice()); buffer } } +impl PartialEq for BufferMut { + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +} + +impl Eq for BufferMut {} + impl Debug for BufferMut { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct(&format!("BufferMut<{}>", type_name::())) @@ -649,7 +817,7 @@ impl BufferMut { let unwritten = self.capacity() - self.len(); // We store `begin` in the case that the lower bound hint is incorrect. - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); // As a first step, we manually iterate the iterator up to the known capacity. @@ -694,7 +862,7 @@ impl BufferMut { .vortex_expect("`TrustedLen` iterator somehow didn't have valid upper bound"), ); - let begin: *const T = self.bytes.spare_capacity_mut().as_mut_ptr().cast(); + let begin: *const T = self.spare_capacity_mut().as_mut_ptr().cast(); let mut dst: *mut T = begin.cast_mut(); iter.for_each(|item| { @@ -799,8 +967,10 @@ impl Buf for ByteBufferMut { self.alignment ); } - self.bytes.advance(cnt); + assert!(cnt <= self.length, "advance out of bounds"); + self.offset += cnt; self.length -= cnt; + self.capacity -= cnt; } } @@ -822,13 +992,15 @@ unsafe impl BufMut for ByteBufferMut { self.alignment ); } - unsafe { self.bytes.advance_mut(cnt) }; - self.length -= cnt; + self.reserve(cnt); + self.length += cnt; } #[inline] fn chunk_mut(&mut self) -> &mut UninitSlice { - self.bytes.chunk_mut() + let spare = self.spare_capacity_mut(); + // SAFETY: spare points to valid uninitialized byte capacity owned by this buffer. + unsafe { UninitSlice::from_raw_parts_mut(spare.as_mut_ptr().cast(), spare.len()) } } fn put(&mut self, mut src: T) @@ -853,35 +1025,6 @@ unsafe impl BufMut for ByteBufferMut { } } -/// Extension trait for [`BytesMut`] that provides functions for aligning the buffer. -trait AlignedBytesMut { - /// Align an empty `BytesMut` to the specified alignment. - /// - /// ## Panics - /// - /// Panics if the buffer is not empty, or if there is not enough capacity to align the buffer. - fn align_empty(&mut self, alignment: Alignment); -} - -impl AlignedBytesMut for BytesMut { - fn align_empty(&mut self, alignment: Alignment) { - // TODO(joe): this is slow fixme - if !self.is_empty() { - vortex_panic!("ByteBufferMut must be empty"); - } - - let padding = self.as_ptr().align_offset(alignment.as_usize()); - self.capacity() - .checked_sub(padding) - .vortex_expect("Not enough capacity to align buffer"); - - // SAFETY: We know the buffer is empty, and we know we have enough capacity, so we can - // safely set the length to the padding and advance the buffer to the aligned offset. - unsafe { self.set_len(padding) }; - self.advance(padding); - } -} - impl Write for ByteBufferMut { fn write(&mut self, buf: &[u8]) -> std::io::Result { self.extend_from_slice(buf); diff --git a/vortex-buffer/src/lib.rs b/vortex-buffer/src/lib.rs index ee113481353..99b1e7c3081 100644 --- a/vortex-buffer/src/lib.rs +++ b/vortex-buffer/src/lib.rs @@ -5,11 +5,10 @@ //! A library for working with custom aligned buffers of sized values. //! -//! The `vortex-buffer` crate is built around `bytes::Bytes` and therefore supports zero-copy -//! cloning and slicing, but differs in that it can define and maintain a custom alignment. +//! The `vortex-buffer` crate supports zero-copy cloning and slicing with a custom allocator and +//! runtime alignment. //! -//! * `Buffer` and `BufferMut` provide immutable and mutable wrappers around `bytes::Bytes` -//! and `bytes::BytesMut` respectively. +//! * `Buffer` and `BufferMut` provide immutable and mutable typed buffers. //! * `ByteBuffer` and `ByteBufferMut` are type aliases for `u8` buffers. //! * `BufferString` is a wrapper around a `ByteBuffer` that enforces utf-8 encoding. //! * `ConstBuffer` provides similar functionality to `Buffer` except with a @@ -47,6 +46,7 @@ //! `arrow_buffer::OffsetBuffer`. pub use alignment::*; +pub use allocation::*; pub use bit::*; pub use buffer::*; pub use buffer_mut::*; @@ -55,6 +55,7 @@ pub use r#const::*; pub use dispatch::*; pub use string::*; mod alignment; +mod allocation; #[cfg(feature = "arrow")] mod arrow; mod bit; diff --git a/vortex-buffer/src/serde.rs b/vortex-buffer/src/serde.rs index 9563236a50a..5d36f8da3d0 100644 --- a/vortex-buffer/src/serde.rs +++ b/vortex-buffer/src/serde.rs @@ -22,7 +22,7 @@ where where S: Serializer, { - serializer.serialize_bytes(self.inner().as_ref()) + serializer.serialize_bytes(self.as_bytes()) } } From 02ac434a7230920336da6940dfe066732961b331 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 27 Aug 2026 17:36:23 -0400 Subject: [PATCH 02/29] fix(buffer): sort allocator dependency Signed-off-by: Nicholas Gates --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 842f4bc38b3..294ed20d996 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -97,6 +97,7 @@ rust-version = "1.95" version = "0.1.0" [workspace.dependencies] +allocator-api2 = "0.2.21" alp = "0.0.2" anyhow = "1.0.100" arbitrary = "1.3.2" @@ -117,7 +118,6 @@ async-fs = "2.2.0" async-lock = "3.4" async-stream = "0.3.6" async-trait = "0.1.89" -allocator-api2 = "0.2.21" base16ct = "1.0.0" bigdecimal = "0.4.8" bindgen = "0.72.0" From 5f30e6369036ec6ea002e8db45b19099dbf82d56 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 10:40:13 -0400 Subject: [PATCH 03/29] fix(buffer): align within raw allocations Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 47 +++++++++++++++------- vortex-buffer/src/buffer_mut.rs | 71 +++++++++++++++++++++++---------- 2 files changed, 83 insertions(+), 35 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 54e62e0d58e..4a58340bb70 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -19,6 +19,8 @@ use crate::Alignment; use crate::BufferMut; /// An allocator that can back a Vortex buffer. +/// +/// Vortex over-allocates raw storage and aligns the buffer within it. pub trait BufferAllocator: Allocator + Debug + Send + Sync + 'static {} impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static {} @@ -190,6 +192,7 @@ pub(crate) struct Allocation { ptr: NonNull, layout: Layout, capacity: usize, + buffer_alignment: Alignment, allocator: BufferAllocatorRef, } @@ -199,20 +202,34 @@ unsafe impl Send for Allocation {} unsafe impl Sync for Allocation {} impl Allocation { - pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { - Self::allocate_impl(layout, allocator, false) - } - - pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { - Self::allocate_impl(layout, allocator, true) - } - - fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { + pub(crate) fn allocate( + layout: Layout, + buffer_alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::allocate_impl(layout, buffer_alignment, allocator, false) + } + + pub(crate) fn allocate_zeroed( + layout: Layout, + buffer_alignment: Alignment, + allocator: BufferAllocatorRef, + ) -> Self { + Self::allocate_impl(layout, buffer_alignment, allocator, true) + } + + fn allocate_impl( + layout: Layout, + buffer_alignment: Alignment, + allocator: BufferAllocatorRef, + zeroed: bool, + ) -> Self { if layout.size() == 0 { return Self { ptr: layout.dangling_ptr(), layout, capacity: 0, + buffer_alignment, allocator, }; } @@ -228,6 +245,7 @@ impl Allocation { ptr: allocation.cast(), layout, capacity: allocation.len(), + buffer_alignment, allocator, } } @@ -240,17 +258,17 @@ impl Allocation { self.capacity } - pub(crate) fn alignment(&self) -> Alignment { - Alignment::new(self.layout.align()) + pub(crate) fn buffer_alignment(&self) -> Alignment { + self.buffer_alignment } pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } - pub(crate) fn grow(&mut self, layout: Layout) { + pub(crate) fn grow(&mut self, layout: Layout, buffer_alignment: Alignment) { if self.layout.size() == 0 { - *self = Self::allocate(layout, self.allocator.clone()); + *self = Self::allocate(layout, buffer_alignment, self.allocator.clone()); return; } @@ -262,6 +280,7 @@ impl Allocation { self.ptr = allocation.cast(); self.layout = layout; self.capacity = allocation.len(); + self.buffer_alignment = buffer_alignment; } } @@ -358,7 +377,7 @@ mod tests { assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!( state.alignment.load(Ordering::Relaxed), - *Alignment::DEFAULT_ALIGNMENT + *Alignment::of::() ); drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 84699ab9695..9bec9dcc337 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -115,14 +115,22 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let layout = Layout::from_size_align(size, actual.as_usize()) + let allocation_size = if size == 0 { + 0 + } else { + size.checked_add(actual.as_usize() - 1) + .vortex_expect("buffer capacity overflow") + }; + let allocation_alignment = if size == 0 { actual.as_usize() } else { 1 }; + let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let allocation = Allocation::allocate(layout, allocator); - let capacity = allocation.size() / size_of::(); + let allocation = Allocation::allocate(layout, actual, allocator); + let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); + let capacity = (allocation.size() - offset) / size_of::(); Self { allocation, - offset: 0, + offset, length: 0, capacity, alignment, @@ -193,13 +201,28 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - let layout = Layout::from_size_align(size, actual_alignment.as_usize()) + let allocation_size = if size == 0 { + 0 + } else { + size.checked_add(actual_alignment.as_usize() - 1) + .vortex_expect("buffer length overflow") + }; + let allocation_alignment = if size == 0 { + actual_alignment.as_usize() + } else { + 1 + }; + let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); - let allocation = Allocation::allocate_zeroed(layout, allocator); - let capacity = allocation.size() / size_of::(); + let allocation = Allocation::allocate_zeroed(layout, actual_alignment, allocator); + let offset = allocation + .ptr() + .as_ptr() + .align_offset(actual_alignment.as_usize()); + let capacity = (allocation.size() - offset) / size_of::(); Self { allocation, - offset: 0, + offset, length: len, capacity, alignment, @@ -444,26 +467,32 @@ impl BufferMut { let new_size = new_capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let physical_alignment = max(self.alignment, self.allocation.alignment()); - let layout = Layout::from_size_align(new_size, physical_alignment.as_usize()) + let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); + let allocation_size = new_size + .checked_add(physical_alignment.as_usize() - 1) + .vortex_expect("buffer capacity overflow"); + let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - if self.offset == 0 { - self.allocation.grow(layout); - } else { - let mut allocation = Allocation::allocate(layout, self.allocation.allocator().clone()); - // SAFETY: both ranges are valid for the initialized byte length and do not overlap. + let old_offset = self.offset; + self.allocation.grow(layout, physical_alignment); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + if old_offset != new_offset { + // SAFETY: both ranges are within the allocation and may overlap. unsafe { - std::ptr::copy_nonoverlapping( - self.as_ptr().cast::(), - allocation.ptr().as_ptr(), + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); } - std::mem::swap(&mut self.allocation, &mut allocation); - self.offset = 0; } - self.capacity = self.allocation.size() / size_of::(); + self.offset = new_offset; + self.capacity = (self.allocation.size() - new_offset) / size_of::(); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. From 3c81101b372f9436ef0f6ae5f7cd9e1dd6e5c9be Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 12:05:39 -0400 Subject: [PATCH 04/29] fix(buffer): avoid realloc when growing buffers Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 46 +++++++++++++++++++++------------ vortex-buffer/src/buffer_mut.rs | 28 ++++++++++---------- 2 files changed, 44 insertions(+), 30 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 4a58340bb70..7ed655e114d 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -265,23 +265,6 @@ impl Allocation { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } - - pub(crate) fn grow(&mut self, layout: Layout, buffer_alignment: Alignment) { - if self.layout.size() == 0 { - *self = Self::allocate(layout, buffer_alignment, self.allocator.clone()); - return; - } - - let allocation = - // SAFETY: ptr and layout describe a live block allocated by self.allocator. The new - // layout is at least as large as the old layout. - unsafe { self.allocator.grow(self.ptr, self.layout, layout) } - .unwrap_or_else(|_| handle_alloc_error(layout)); - self.ptr = allocation.cast(); - self.layout = layout; - self.capacity = allocation.len(); - self.buffer_alignment = buffer_alignment; - } } impl Drop for Allocation { @@ -345,6 +328,7 @@ mod tests { struct TrackingState { allocations: AtomicUsize, deallocations: AtomicUsize, + grows: AtomicUsize, alignment: AtomicUsize, } @@ -363,6 +347,17 @@ mod tests { // SAFETY: the caller passes the pointer and layout returned by Global. unsafe { Global.deallocate(ptr, layout) } } + + unsafe fn grow( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + self.state.grows.fetch_add(1, Ordering::Relaxed); + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } } #[test] @@ -384,4 +379,21 @@ mod tests { drop(view); assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); } + + #[test] + fn buffer_growth_allocates_and_copies() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(1); + let initial_capacity = buffer.capacity(); + buffer.extend(std::iter::repeat_n(7, initial_capacity)); + + buffer.push(u32::MAX); + + assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]); + assert_eq!(buffer[initial_capacity], u32::MAX); + assert_eq!(state.allocations.load(Ordering::Relaxed), 2); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + } } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 9bec9dcc337..1477b3f2026 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -474,23 +474,25 @@ impl BufferMut { let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let old_offset = self.offset; - self.allocation.grow(layout, physical_alignment); - let new_offset = self - .allocation + let allocation = Allocation::allocate( + layout, + physical_alignment, + self.allocation.allocator().clone(), + ); + let new_offset = allocation .ptr() .as_ptr() .align_offset(physical_alignment.as_usize()); - if old_offset != new_offset { - // SAFETY: both ranges are within the allocation and may overlap. - unsafe { - std::ptr::copy( - self.allocation.ptr().as_ptr().add(old_offset), - self.allocation.ptr().as_ptr().add(new_offset), - self.length * size_of::(), - ); - } + // SAFETY: the source contains `length` initialized elements and the fresh allocation has + // room for at least `new_capacity` elements. The allocations do not overlap. + unsafe { + std::ptr::copy_nonoverlapping( + self.allocation.ptr().as_ptr().add(self.offset), + allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); } + self.allocation = allocation; self.offset = new_offset; self.capacity = (self.allocation.size() - new_offset) / size_of::(); } From 758cbe22fb262513799a90a147da63b8c3f47d03 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 13:02:20 -0400 Subject: [PATCH 05/29] perf(buffer): avoid indirection for static allocator Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 84 +++++++++++++++++++++++++-------- 1 file changed, 65 insertions(+), 19 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 7ed655e114d..2519241e4a6 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -8,7 +8,6 @@ use std::fmt; use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; -use std::sync::LazyLock; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; @@ -27,17 +26,18 @@ impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static /// A shared reference to a buffer allocator. #[derive(Clone)] -pub struct BufferAllocatorRef(Arc); +pub struct BufferAllocatorRef(Option>); impl BufferAllocatorRef { /// Wrap an allocator in a shared reference. pub fn new(allocator: impl BufferAllocator) -> Self { - Self(Arc::new(allocator)) + Self(Some(Arc::new(allocator))) } /// Return a shared reference to the static allocator. + #[inline] pub fn statically_allocated() -> Self { - STATIC_ALLOCATOR.clone() + Self(None) } /// Create a mutable buffer with this allocator. @@ -63,53 +63,100 @@ impl BufferAllocatorRef { impl Debug for BufferAllocatorRef { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt(f) + match &self.0 { + Some(allocator) => allocator.fmt(f), + None => StaticBufferAllocator.fmt(f), + } } } -// SAFETY: all calls are forwarded to the same allocator value held by the Arc. +// SAFETY: all calls are forwarded to the allocator represented by this value. unsafe impl Allocator for BufferAllocatorRef { + #[inline] fn allocate(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate(layout) + match &self.0 { + Some(allocator) => allocator.allocate(layout), + None => Global.allocate(layout), + } } + #[inline] fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate_zeroed(layout) + match &self.0 { + Some(allocator) => allocator.allocate_zeroed(layout), + None => Global.allocate_zeroed(layout), + } } + #[inline] unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.deallocate(ptr, layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.deallocate(ptr, layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.deallocate(ptr, layout) } + } + } } + #[inline] unsafe fn grow( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.grow(ptr, old_layout, new_layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow(ptr, old_layout, new_layout) } + } + } } + #[inline] unsafe fn grow_zeroed( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.grow_zeroed(ptr, old_layout, new_layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } + } + } } + #[inline] unsafe fn shrink( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.shrink(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { allocator.shrink(ptr, old_layout, new_layout) } + } + None => { + // SAFETY: the caller upholds the Allocator contract. + unsafe { Global.shrink(ptr, old_layout, new_layout) } + } + } } } @@ -185,8 +232,7 @@ unsafe impl Allocator for StaticBufferAllocator { } } -static STATIC_ALLOCATOR: LazyLock = - LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); +static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); pub(crate) struct Allocation { ptr: NonNull, @@ -299,7 +345,7 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), + Self::External { .. } => &STATIC_ALLOCATOR, } } } From 03b068d1aefa1b086c3a57774fe282a877595211 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 13:33:23 -0400 Subject: [PATCH 06/29] Revert "perf(buffer): avoid indirection for static allocator" This reverts commit 1186945ab2f8c0b4d9f6e6dbbcdcf82cf67ac40b. Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 84 ++++++++------------------------- 1 file changed, 19 insertions(+), 65 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 2519241e4a6..7ed655e114d 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -8,6 +8,7 @@ use std::fmt; use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; +use std::sync::LazyLock; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; @@ -26,18 +27,17 @@ impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static /// A shared reference to a buffer allocator. #[derive(Clone)] -pub struct BufferAllocatorRef(Option>); +pub struct BufferAllocatorRef(Arc); impl BufferAllocatorRef { /// Wrap an allocator in a shared reference. pub fn new(allocator: impl BufferAllocator) -> Self { - Self(Some(Arc::new(allocator))) + Self(Arc::new(allocator)) } /// Return a shared reference to the static allocator. - #[inline] pub fn statically_allocated() -> Self { - Self(None) + STATIC_ALLOCATOR.clone() } /// Create a mutable buffer with this allocator. @@ -63,100 +63,53 @@ impl BufferAllocatorRef { impl Debug for BufferAllocatorRef { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match &self.0 { - Some(allocator) => allocator.fmt(f), - None => StaticBufferAllocator.fmt(f), - } + self.0.fmt(f) } } -// SAFETY: all calls are forwarded to the allocator represented by this value. +// SAFETY: all calls are forwarded to the same allocator value held by the Arc. unsafe impl Allocator for BufferAllocatorRef { - #[inline] fn allocate(&self, layout: Layout) -> Result, AllocError> { - match &self.0 { - Some(allocator) => allocator.allocate(layout), - None => Global.allocate(layout), - } + self.0.allocate(layout) } - #[inline] fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - match &self.0 { - Some(allocator) => allocator.allocate_zeroed(layout), - None => Global.allocate_zeroed(layout), - } + self.0.allocate_zeroed(layout) } - #[inline] unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.deallocate(ptr, layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.deallocate(ptr, layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.deallocate(ptr, layout) } } - #[inline] unsafe fn grow( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.grow(ptr, old_layout, new_layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.grow(ptr, old_layout, new_layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow(ptr, old_layout, new_layout) } } - #[inline] unsafe fn grow_zeroed( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.grow_zeroed(ptr, old_layout, new_layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } } - #[inline] unsafe fn shrink( &self, ptr: NonNull, old_layout: Layout, new_layout: Layout, ) -> Result, AllocError> { - match &self.0 { - Some(allocator) => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { allocator.shrink(ptr, old_layout, new_layout) } - } - None => { - // SAFETY: the caller upholds the Allocator contract. - unsafe { Global.shrink(ptr, old_layout, new_layout) } - } - } + // SAFETY: the caller upholds the Allocator contract. + unsafe { self.0.shrink(ptr, old_layout, new_layout) } } } @@ -232,7 +185,8 @@ unsafe impl Allocator for StaticBufferAllocator { } } -static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); +static STATIC_ALLOCATOR: LazyLock = + LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); pub(crate) struct Allocation { ptr: NonNull, @@ -345,7 +299,7 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::External { .. } => &STATIC_ALLOCATOR, + Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), } } } From 39c0f8801b6ef1b24ed05b59573e0f87c11d1154 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 13:35:05 -0400 Subject: [PATCH 07/29] perf(buffer): inline hot buffer growth paths Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 5 +++++ vortex-buffer/src/buffer_mut.rs | 6 ++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 7ed655e114d..c1eb4a32139 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -250,18 +250,22 @@ impl Allocation { } } + #[inline(always)] pub(crate) fn ptr(&self) -> NonNull { self.ptr } + #[inline(always)] pub(crate) fn size(&self) -> usize { self.capacity } + #[inline(always)] pub(crate) fn buffer_alignment(&self) -> Alignment { self.buffer_alignment } + #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } @@ -296,6 +300,7 @@ pub(crate) enum BufferBacking { } impl BufferBacking { + #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 1477b3f2026..c858e44bdc1 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -398,12 +398,14 @@ impl BufferMut { } /// Returns a raw pointer to the buffer's data. + #[inline(always)] pub fn as_ptr(&self) -> *const T { // SAFETY: offset always remains within the allocation. unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } } /// Returns a mutable raw pointer to the buffer's data. + #[inline(always)] pub fn as_mut_ptr(&mut self) -> *mut T { // SAFETY: BufferMut uniquely owns the allocation and offset is in bounds. unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } @@ -974,8 +976,8 @@ where impl FromIterator for BufferMut { fn from_iter>(iter: I) -> Self { - // We don't infer the capacity here and just let the first call to `extend` do it for us. - let mut buffer = Self::with_capacity(0); + let iter = iter.into_iter(); + let mut buffer = Self::with_capacity(iter.size_hint().0); buffer.extend(iter); buffer } From 7e3e8264473ebf1a24e5a7c975500ce115153f20 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 14:02:12 -0400 Subject: [PATCH 08/29] perf(buffer): preserve aligned seed capacity Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 28 ++++++++-------------------- 1 file changed, 8 insertions(+), 20 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index c858e44bdc1..26de935a03d 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -115,14 +115,10 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let allocation_size = if size == 0 { - 0 - } else { - size.checked_add(actual.as_usize() - 1) - .vortex_expect("buffer capacity overflow") - }; - let allocation_alignment = if size == 0 { actual.as_usize() } else { 1 }; - let layout = Layout::from_size_align(allocation_size, allocation_alignment) + let allocation_size = size + .checked_add(actual.as_usize()) + .vortex_expect("buffer capacity overflow"); + let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let allocation = Allocation::allocate(layout, actual, allocator); let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); @@ -201,18 +197,10 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - let allocation_size = if size == 0 { - 0 - } else { - size.checked_add(actual_alignment.as_usize() - 1) - .vortex_expect("buffer length overflow") - }; - let allocation_alignment = if size == 0 { - actual_alignment.as_usize() - } else { - 1 - }; - let layout = Layout::from_size_align(allocation_size, allocation_alignment) + let allocation_size = size + .checked_add(actual_alignment.as_usize()) + .vortex_expect("buffer length overflow"); + let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); let allocation = Allocation::allocate_zeroed(layout, actual_alignment, allocator); let offset = allocation From 1d6465693363681e6e4cf000f572fbe4a86e48bc Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 14:22:41 -0400 Subject: [PATCH 09/29] perf(buffer): restore byte-based growth Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 26de935a03d..7864d909e23 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -452,15 +452,15 @@ impl BufferMut { .length .checked_add(additional) .vortex_expect("buffer capacity overflow"); - // Make sure we at least double in size each time we re-allocate to amortize the cost - let new_capacity = required.max(self.capacity.saturating_mul(2)); - let new_size = new_capacity + let required_size = required .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); - let allocation_size = new_size - .checked_add(physical_alignment.as_usize() - 1) + let required_size = required_size + .checked_add(physical_alignment.as_usize()) .vortex_expect("buffer capacity overflow"); + let current_size = self.allocation.size() - self.offset; + let allocation_size = required_size.max(current_size.saturating_mul(2)); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); @@ -474,7 +474,7 @@ impl BufferMut { .as_ptr() .align_offset(physical_alignment.as_usize()); // SAFETY: the source contains `length` initialized elements and the fresh allocation has - // room for at least `new_capacity` elements. The allocations do not overlap. + // room for at least `required` elements. The allocations do not overlap. unsafe { std::ptr::copy_nonoverlapping( self.allocation.ptr().as_ptr().add(self.offset), From e921aa6bfaac1f299e47b3e05f191a74b959b173 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 16:40:52 -0400 Subject: [PATCH 10/29] perf(buffer): compact allocator-backed storage Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 144 ++++++++++++++++++++----------- vortex-buffer/src/bit/buf_mut.rs | 19 ---- vortex-buffer/src/buffer.rs | 101 +++++++++++++++++----- vortex-buffer/src/buffer_mut.rs | 134 +++++++++++----------------- 4 files changed, 224 insertions(+), 174 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index c1eb4a32139..5ba2e54b44c 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -6,14 +6,15 @@ use std::alloc::Layout; use std::fmt; use std::fmt::Debug; +use std::mem::ManuallyDrop; use std::ptr::NonNull; use std::sync::Arc; -use std::sync::LazyLock; use allocator_api2::alloc::AllocError; use allocator_api2::alloc::Allocator; use allocator_api2::alloc::Global; use allocator_api2::alloc::handle_alloc_error; +use vortex_error::VortexExpect; use crate::Alignment; use crate::BufferMut; @@ -27,17 +28,17 @@ impl BufferAllocator for A where A: Allocator + Debug + Send + Sync + 'static /// A shared reference to a buffer allocator. #[derive(Clone)] -pub struct BufferAllocatorRef(Arc); +pub struct BufferAllocatorRef(Option>); impl BufferAllocatorRef { /// Wrap an allocator in a shared reference. pub fn new(allocator: impl BufferAllocator) -> Self { - Self(Arc::new(allocator)) + Self(Some(Arc::new(allocator))) } /// Return a shared reference to the static allocator. pub fn statically_allocated() -> Self { - STATIC_ALLOCATOR.clone() + Self(None) } /// Create a mutable buffer with this allocator. @@ -63,23 +64,35 @@ impl BufferAllocatorRef { impl Debug for BufferAllocatorRef { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - self.0.fmt(f) + match &self.0 { + Some(allocator) => allocator.fmt(f), + None => StaticBufferAllocator.fmt(f), + } } } // SAFETY: all calls are forwarded to the same allocator value held by the Arc. unsafe impl Allocator for BufferAllocatorRef { fn allocate(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate(layout) + match &self.0 { + Some(allocator) => allocator.allocate(layout), + None => Global.allocate(layout), + } } fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { - self.0.allocate_zeroed(layout) + match &self.0 { + Some(allocator) => allocator.allocate_zeroed(layout), + None => Global.allocate_zeroed(layout), + } } unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.deallocate(ptr, layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.deallocate(ptr, layout) }, + None => unsafe { Global.deallocate(ptr, layout) }, + } } unsafe fn grow( @@ -89,7 +102,10 @@ unsafe impl Allocator for BufferAllocatorRef { new_layout: Layout, ) -> Result, AllocError> { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.grow(ptr, old_layout, new_layout) }, + None => unsafe { Global.grow(ptr, old_layout, new_layout) }, + } } unsafe fn grow_zeroed( @@ -99,7 +115,10 @@ unsafe impl Allocator for BufferAllocatorRef { new_layout: Layout, ) -> Result, AllocError> { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.grow_zeroed(ptr, old_layout, new_layout) }, + None => unsafe { Global.grow_zeroed(ptr, old_layout, new_layout) }, + } } unsafe fn shrink( @@ -109,7 +128,10 @@ unsafe impl Allocator for BufferAllocatorRef { new_layout: Layout, ) -> Result, AllocError> { // SAFETY: the caller upholds the Allocator contract. - unsafe { self.0.shrink(ptr, old_layout, new_layout) } + match &self.0 { + Some(allocator) => unsafe { allocator.shrink(ptr, old_layout, new_layout) }, + None => unsafe { Global.shrink(ptr, old_layout, new_layout) }, + } } } @@ -185,14 +207,11 @@ unsafe impl Allocator for StaticBufferAllocator { } } -static STATIC_ALLOCATOR: LazyLock = - LazyLock::new(|| BufferAllocatorRef::new(StaticBufferAllocator)); +static STATIC_ALLOCATOR: BufferAllocatorRef = BufferAllocatorRef(None); pub(crate) struct Allocation { ptr: NonNull, layout: Layout, - capacity: usize, - buffer_alignment: Alignment, allocator: BufferAllocatorRef, } @@ -202,34 +221,35 @@ unsafe impl Send for Allocation {} unsafe impl Sync for Allocation {} impl Allocation { - pub(crate) fn allocate( - layout: Layout, - buffer_alignment: Alignment, - allocator: BufferAllocatorRef, - ) -> Self { - Self::allocate_impl(layout, buffer_alignment, allocator, false) - } - - pub(crate) fn allocate_zeroed( - layout: Layout, - buffer_alignment: Alignment, - allocator: BufferAllocatorRef, - ) -> Self { - Self::allocate_impl(layout, buffer_alignment, allocator, true) - } - - fn allocate_impl( - layout: Layout, - buffer_alignment: Alignment, - allocator: BufferAllocatorRef, - zeroed: bool, - ) -> Self { + pub(crate) fn allocate(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, false) + } + + pub(crate) fn allocate_zeroed(layout: Layout, allocator: BufferAllocatorRef) -> Self { + Self::allocate_impl(layout, allocator, true) + } + + pub(crate) fn from_vec(vec: Vec) -> Self { + assert!(!std::mem::needs_drop::()); + + let mut vec = ManuallyDrop::new(vec); + let layout = Layout::array::(vec.capacity()) + .unwrap_or_else(|_| unreachable!("a Vec capacity always has a valid layout")); + let ptr = NonNull::new(vec.as_mut_ptr().cast()) + .vortex_expect("a Vec always has a non-null pointer"); + + Self { + ptr, + layout, + allocator: BufferAllocatorRef::statically_allocated(), + } + } + + fn allocate_impl(layout: Layout, allocator: BufferAllocatorRef, zeroed: bool) -> Self { if layout.size() == 0 { return Self { ptr: layout.dangling_ptr(), layout, - capacity: 0, - buffer_alignment, allocator, }; } @@ -244,8 +264,6 @@ impl Allocation { Self { ptr: allocation.cast(), layout, - capacity: allocation.len(), - buffer_alignment, allocator, } } @@ -257,18 +275,31 @@ impl Allocation { #[inline(always)] pub(crate) fn size(&self) -> usize { - self.capacity + self.layout.size() } #[inline(always)] - pub(crate) fn buffer_alignment(&self) -> Alignment { - self.buffer_alignment + pub(crate) fn alignment(&self) -> usize { + self.layout.align() } #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator } + + pub(crate) fn grow(&mut self, new_layout: Layout) { + let allocation = if self.layout.size() == 0 { + self.allocator.allocate(new_layout) + } else { + // SAFETY: ptr denotes a live block owned by allocator, old_layout fits the block, and + // the caller only grows the allocation. + unsafe { self.allocator.grow(self.ptr, self.layout, new_layout) } + } + .unwrap_or_else(|_| handle_alloc_error(new_layout)); + self.ptr = allocation.cast(); + self.layout = new_layout; + } } impl Drop for Allocation { @@ -282,15 +313,21 @@ impl Drop for Allocation { } pub(crate) trait BufferOwner: Send + Sync + 'static { - fn as_slice(&self) -> &[u8]; + fn as_ptr(&self) -> *const u8; + + fn len(&self) -> usize; } impl BufferOwner for T where T: AsRef<[u8]> + Send + Sync + 'static, { - fn as_slice(&self) -> &[u8] { - self.as_ref() + fn as_ptr(&self) -> *const u8 { + self.as_ref().as_ptr() + } + + fn len(&self) -> usize { + self.as_ref().len() } } @@ -304,7 +341,7 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::External { .. } => LazyLock::force(&STATIC_ALLOCATOR), + Self::External { .. } => &STATIC_ALLOCATOR, } } } @@ -377,7 +414,7 @@ mod tests { assert_eq!(state.allocations.load(Ordering::Relaxed), 1); assert_eq!( state.alignment.load(Ordering::Relaxed), - *Alignment::of::() + Alignment::of::().as_usize() ); drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); @@ -386,7 +423,7 @@ mod tests { } #[test] - fn buffer_growth_allocates_and_copies() { + fn buffer_growth_uses_allocator_grow() { let allocator = TrackingAllocator::default(); let state = Arc::clone(&allocator.state); let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(1); @@ -397,8 +434,11 @@ mod tests { assert_eq!(&buffer[..initial_capacity], vec![7; initial_capacity]); assert_eq!(buffer[initial_capacity], u32::MAX); - assert_eq!(state.allocations.load(Ordering::Relaxed), 2); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + assert_eq!(state.grows.load(Ordering::Relaxed), 1); + + drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); - assert_eq!(state.grows.load(Ordering::Relaxed), 0); } } diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index 99935f627f3..2e5c3ea39b6 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -588,25 +588,6 @@ impl BitBufferMut { self.len += bit_len; } - /// Absorbs a mutable buffer that was previously split off. - /// - /// If the two buffers were previously contiguous and not mutated in a way that causes - /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is - /// an O(1) operation that just decreases a reference count and sets a few indices. - /// - /// Otherwise, this method degenerates to self.append_buffer(&other). - pub fn unsplit(&mut self, other: Self) { - if (self.offset + self.len).is_multiple_of(8) && other.offset == 0 { - // We are aligned and can just append the buffers - self.buffer.unsplit(other.buffer); - self.len += other.len; - return; - } - - // Otherwise, we need to append the bits one by one - self.append_buffer(&other.freeze()) - } - /// Freeze the buffer in its current state into an immutable `BoolBuffer`. #[inline] pub fn freeze(self) -> BitBuffer { diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 4b5ee8ef6fd..026d3bcf492 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -34,6 +34,7 @@ pub struct Buffer { pub(crate) ptr: NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, + pub(crate) physical_alignment: Alignment, pub(crate) backing: Arc, } @@ -64,6 +65,7 @@ impl Default for Buffer { ptr: empty_ptr(), length: 0, alignment: Alignment::of::(), + physical_alignment: Alignment::MAX, backing: EMPTY_BACKING.clone(), } } @@ -105,6 +107,7 @@ impl Buffer { offset: usize, length: usize, alignment: Alignment, + physical_alignment: Alignment, ) -> Self { // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. let ptr = unsafe { allocation.ptr().add(offset).cast() }; @@ -112,25 +115,25 @@ impl Buffer { ptr, length, alignment, + physical_alignment, backing: Arc::new(BufferBacking::Owned(allocation)), } } fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self { - let bytes = owner.as_slice(); - let length = bytes.len() / size_of::(); + let owner: Box = Box::new(owner); + let length = owner.len() / size_of::(); let ptr = if length == 0 { empty_ptr() } else { - NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") + NonNull::new(owner.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") }; Self { ptr, length, alignment, - backing: Arc::new(BufferBacking::External { - _owner: Box::new(owner), - }), + physical_alignment: alignment, + backing: Arc::new(BufferBacking::External { _owner: owner }), } } @@ -224,6 +227,7 @@ impl Buffer { ptr: empty_ptr(), length: 0, alignment, + physical_alignment: Alignment::MAX, backing: EMPTY_BACKING.clone(), } } @@ -283,6 +287,7 @@ impl Buffer { ptr: buffer.ptr.cast(), length: buffer.length / size_of::(), alignment, + physical_alignment: buffer.physical_alignment, backing: buffer.backing, } } @@ -494,6 +499,7 @@ impl Buffer { ptr: unsafe { self.ptr.add(begin) }, length: end - begin, alignment, + physical_alignment: self.physical_alignment, backing: Arc::clone(&self.backing), } } @@ -548,6 +554,7 @@ impl Buffer { ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"), length: subset.len(), alignment, + physical_alignment: self.physical_alignment, backing: Arc::clone(&self.backing), } } @@ -567,6 +574,7 @@ impl Buffer { ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, + physical_alignment: self.physical_alignment, backing: self.backing, } } @@ -577,18 +585,18 @@ impl Buffer { ptr, length, alignment, + physical_alignment, backing, } = self; match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); - let capacity = (allocation.size() - offset) / size_of::(); Ok(BufferMut { allocation, offset, length, - capacity, alignment, + physical_alignment, _marker: Default::default(), }) } @@ -596,12 +604,14 @@ impl Buffer { ptr, length, alignment, + physical_alignment, backing: Arc::new(backing), }), Err(backing) => Err(Self { ptr, length, alignment, + physical_alignment, backing, }), } @@ -673,6 +683,7 @@ impl Buffer { ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, + physical_alignment: self.physical_alignment, backing: self.backing, } } @@ -754,15 +765,17 @@ impl FromIterator for Buffer { } } -// Helper struct to allow us to zero-copy any vec into a buffer +// Helper struct that preserves drop glue for non-native Vec elements. #[repr(transparent)] struct Wrapper(Vec); -impl AsRef<[u8]> for Wrapper { - fn as_ref(&self) -> &[u8] { - let data = self.0.as_ptr().cast::(); - let len = self.0.len() * size_of::(); - unsafe { std::slice::from_raw_parts(data, len) } +impl crate::BufferOwner for Wrapper { + fn as_ptr(&self) -> *const u8 { + self.0.as_ptr().cast() + } + + fn len(&self) -> usize { + self.0.len() * size_of::() } } @@ -771,12 +784,13 @@ where T: Send + Sync + 'static, { fn from(value: Vec) -> Self { - let wrapped_vec = Wrapper(value); - assert_eq!( - wrapped_vec.as_ref().as_ptr().align_offset(align_of::()), - 0 - ); - Self::from_owner(wrapped_vec, Alignment::of::()) + let length = value.len(); + let alignment = Alignment::of::(); + if std::mem::needs_drop::() { + Self::from_owner(Wrapper(value), alignment) + } else { + Self::from_allocation(Allocation::from_vec(value), 0, length, alignment, alignment) + } } } @@ -898,6 +912,11 @@ impl From> for Buffer { #[cfg(test)] mod test { + use std::mem::align_of; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use bytes::Buf; use crate::Alignment; @@ -1015,6 +1034,48 @@ mod test { assert_eq!(vec, buff.as_ref()); } + #[test] + fn from_vec_adopts_allocation() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u32, 2, 3, 4, 5]); + let ptr = vec.as_ptr(); + let capacity = vec.capacity(); + + let buffer = Buffer::from(vec); + assert_eq!(buffer.as_ptr(), ptr); + + let Ok(mut buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), capacity); + assert_eq!(buffer.allocation.alignment(), align_of::()); + + buffer.extend(6..=32); + assert_eq!(buffer.as_slice(), (1..=32).collect::>()); + assert_eq!(buffer.allocation.alignment(), align_of::()); + } + + #[test] + fn from_vec_preserves_drop_glue() { + struct DropValue(Arc); + + impl Drop for DropValue { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + let drops = Arc::new(AtomicUsize::new(0)); + let values = (0..3) + .map(|_| DropValue(Arc::clone(&drops))) + .collect::>(); + let buffer = Buffer::from(values); + + assert_eq!(drops.load(Ordering::Relaxed), 0); + drop(buffer); + assert_eq!(drops.load(Ordering::Relaxed), 3); + } + #[test] fn empty_aligned_max_alignment() { // Empty buffers are backed by a static and must satisfy any valid alignment. diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 7864d909e23..930d8a1e0f7 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -31,8 +31,8 @@ pub struct BufferMut { pub(crate) allocation: Allocation, pub(crate) offset: usize, pub(crate) length: usize, - pub(crate) capacity: usize, pub(crate) alignment: Alignment, + pub(crate) physical_alignment: Alignment, pub(crate) _marker: std::marker::PhantomData, } @@ -120,16 +120,14 @@ impl BufferMut { .vortex_expect("buffer capacity overflow"); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let allocation = Allocation::allocate(layout, actual, allocator); + let allocation = Allocation::allocate(layout, allocator); let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); - let capacity = (allocation.size() - offset) / size_of::(); - Self { allocation, offset, length: 0, - capacity, alignment, + physical_alignment: actual, _marker: Default::default(), } } @@ -202,18 +200,17 @@ impl BufferMut { .vortex_expect("buffer length overflow"); let layout = Layout::from_size_align(allocation_size, 1) .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); - let allocation = Allocation::allocate_zeroed(layout, actual_alignment, allocator); + let allocation = Allocation::allocate_zeroed(layout, allocator); let offset = allocation .ptr() .as_ptr() .align_offset(actual_alignment.as_usize()); - let capacity = (allocation.size() - offset) / size_of::(); Self { allocation, offset, length: len, - capacity, alignment, + physical_alignment: actual_alignment, _marker: Default::default(), } } @@ -382,7 +379,7 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - self.capacity + (self.allocation.size() - self.offset) / size_of::() } /// Returns a raw pointer to the buffer's data. @@ -437,7 +434,7 @@ impl BufferMut { /// Reserves capacity for at least `additional` more elements to be inserted in the buffer. #[inline] pub fn reserve(&mut self, additional: usize) { - if additional <= self.capacity - self.length { + if additional <= self.capacity() - self.length { // We can fit the additional bytes in the remaining capacity. Nothing to do. return; } @@ -455,36 +452,35 @@ impl BufferMut { let required_size = required .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let physical_alignment = max(self.alignment, self.allocation.buffer_alignment()); + let physical_alignment = max(self.alignment, self.physical_alignment); let required_size = required_size .checked_add(physical_alignment.as_usize()) .vortex_expect("buffer capacity overflow"); let current_size = self.allocation.size() - self.offset; let allocation_size = required_size.max(current_size.saturating_mul(2)); - let layout = Layout::from_size_align(allocation_size, 1) + let layout = Layout::from_size_align(allocation_size, self.allocation.alignment()) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let allocation = Allocation::allocate( - layout, - physical_alignment, - self.allocation.allocator().clone(), - ); - let new_offset = allocation + let old_offset = self.offset; + self.allocation.grow(layout); + let new_offset = self + .allocation .ptr() .as_ptr() .align_offset(physical_alignment.as_usize()); - // SAFETY: the source contains `length` initialized elements and the fresh allocation has - // room for at least `required` elements. The allocations do not overlap. - unsafe { - std::ptr::copy_nonoverlapping( - self.allocation.ptr().as_ptr().add(self.offset), - allocation.ptr().as_ptr().add(new_offset), - self.length * size_of::(), - ); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits overlap. + unsafe { + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } } - self.allocation = allocation; self.offset = new_offset; - self.capacity = (self.allocation.size() - new_offset) / size_of::(); + self.physical_alignment = physical_alignment; } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -626,71 +622,27 @@ impl BufferMut { self.length += slice.len(); } - /// Splits the buffer into two at the given index. - /// - /// Afterward, self contains elements `[0, at)`, and the returned buffer contains elements - /// `[at, capacity)`. The returned buffer uses a new allocation. - /// - /// Panics if either half would have a length that is not a multiple of the alignment. - pub fn split_off(&mut self, at: usize) -> Self { - if at > self.capacity() { - vortex_panic!("Cannot split buffer of capacity {} at {}", self.len(), at); - } - - let bytes_at = at * size_of::(); - if !self.alignment.is_offset_aligned(bytes_at) { - vortex_panic!( - "Cannot split buffer at {}, resulting alignment is not {}", - at, - self.alignment - ); - } - - let new_length = self.length.saturating_sub(at); - let new_capacity = self.capacity - at; - let mut other = Self::with_capacity_aligned_in( - new_capacity, - self.alignment, - self.allocation.allocator().clone(), - ); - if new_length > 0 { - other.extend_from_slice(&self.as_slice()[at..]); - } - self.length = self.length.min(at); - self.capacity = at; - - other - } - - /// Absorbs a mutable buffer that was previously split off. - /// - /// This appends the contents of `other` to this buffer. - pub fn unsplit(&mut self, other: Self) { - if self.alignment != other.alignment { - vortex_panic!( - "Cannot unsplit buffers with different alignments: {} and {}", - self.alignment, - other.alignment - ); - } - self.extend_from_slice(other.as_slice()); - } - /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { ByteBufferMut { allocation: self.allocation, offset: self.offset, length: self.length * size_of::(), - capacity: self.capacity * size_of::(), alignment: self.alignment, + physical_alignment: self.physical_alignment, _marker: Default::default(), } } /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { - Buffer::from_allocation(self.allocation, self.offset, self.length, self.alignment) + Buffer::from_allocation( + self.allocation, + self.offset, + self.length, + self.alignment, + self.physical_alignment, + ) } /// Map each element of the buffer with a closure. @@ -718,7 +670,11 @@ impl BufferMut { /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { if self.as_ptr().align_offset(alignment.as_usize()) == 0 { - Self { alignment, ..self } + Self { + alignment, + physical_alignment: max(self.physical_alignment, alignment), + ..self + } } else { let allocator = self.allocation.allocator().clone(); Self::copy_from_aligned_in(self, alignment, allocator) @@ -748,8 +704,8 @@ impl BufferMut { allocation: self.allocation, offset: self.offset, length: self.length, - capacity: self.capacity, alignment: self.alignment, + physical_alignment: self.physical_alignment, _marker: std::marker::PhantomData, } } @@ -991,7 +947,6 @@ impl Buf for ByteBufferMut { assert!(cnt <= self.length, "advance out of bounds"); self.offset += cnt; self.length -= cnt; - self.capacity -= cnt; } } @@ -1082,6 +1037,19 @@ mod test { assert_eq!(buf.alignment(), Alignment::new(1024)); } + #[test] + fn growth_preserves_alignment_and_values() { + let alignment = Alignment::new(4096); + let mut buffer = BufferMut::::with_capacity_aligned(1, alignment); + + for value in 0..10_000 { + buffer.push(value); + assert!(alignment.is_offset_aligned(buffer.as_ptr().addr())); + } + + assert_eq!(buffer.as_slice(), (0..10_000).collect::>()); + } + #[test] fn from_iter() { let buf = BufferMut::from_iter([0, 10, 20, 30]); From 5447b4352735671d4b7083c93824aa7b7751ae81 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 20:05:57 -0400 Subject: [PATCH 11/29] perf(buffer): avoid empty data allocations Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 17 +++++++++++++++ vortex-buffer/src/buffer_mut.rs | 38 +++++++++++++++++++++++---------- 2 files changed, 44 insertions(+), 11 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 5ba2e54b44c..107c7271999 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -441,4 +441,21 @@ mod tests { drop(buffer); assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); } + + #[test] + fn zero_capacity_does_not_allocate() { + let allocator = TrackingAllocator::default(); + let state = Arc::clone(&allocator.state); + let mut buffer = BufferAllocatorRef::new(allocator).with_capacity::(0); + + assert_eq!(buffer.capacity(), 0); + assert!(Alignment::DEFAULT_ALIGNMENT.is_offset_aligned(buffer.as_ptr().addr())); + assert_eq!(state.allocations.load(Ordering::Relaxed), 0); + + buffer.push(42); + + assert_eq!(buffer.as_slice(), [42]); + assert_eq!(state.allocations.load(Ordering::Relaxed), 1); + assert_eq!(state.grows.load(Ordering::Relaxed), 0); + } } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 930d8a1e0f7..26fc088ff2a 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -115,11 +115,17 @@ impl BufferMut { let size = capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let allocation_size = size - .checked_add(actual.as_usize()) - .vortex_expect("buffer capacity overflow"); - let layout = Layout::from_size_align(allocation_size, 1) - .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); + let layout = if size == 0 { + Layout::from_size_align(0, actual.as_usize()) + .unwrap_or_else(|_| vortex_panic!("invalid empty buffer alignment")) + } else { + let allocation_size = size + .checked_add(actual.as_usize()) + .vortex_expect("buffer capacity overflow"); + Layout::from_size_align(allocation_size, 1).unwrap_or_else(|_| { + vortex_panic!("buffer capacity exceeds maximum allocation size") + }) + }; let allocation = Allocation::allocate(layout, allocator); let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); Self { @@ -195,11 +201,16 @@ impl BufferMut { let size = len .checked_mul(size_of::()) .vortex_expect("buffer length overflow"); - let allocation_size = size - .checked_add(actual_alignment.as_usize()) - .vortex_expect("buffer length overflow"); - let layout = Layout::from_size_align(allocation_size, 1) - .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")); + let layout = if size == 0 { + Layout::from_size_align(0, actual_alignment.as_usize()) + .unwrap_or_else(|_| vortex_panic!("invalid empty buffer alignment")) + } else { + let allocation_size = size + .checked_add(actual_alignment.as_usize()) + .vortex_expect("buffer length overflow"); + Layout::from_size_align(allocation_size, 1) + .unwrap_or_else(|_| vortex_panic!("buffer length exceeds maximum allocation size")) + }; let allocation = Allocation::allocate_zeroed(layout, allocator); let offset = allocation .ptr() @@ -458,7 +469,12 @@ impl BufferMut { .vortex_expect("buffer capacity overflow"); let current_size = self.allocation.size() - self.offset; let allocation_size = required_size.max(current_size.saturating_mul(2)); - let layout = Layout::from_size_align(allocation_size, self.allocation.alignment()) + let allocation_alignment = if self.allocation.size() == 0 { + 1 + } else { + self.allocation.alignment() + }; + let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let old_offset = self.offset; From e342a1dae226c5e064f1314da00f2236c35f1296 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 21:25:23 -0400 Subject: [PATCH 12/29] perf(buffer): copy live data for static growth Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 4 +++ vortex-buffer/src/buffer_mut.rs | 45 ++++++++++++++++++++++++--------- 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 107c7271999..415ba4fc4b4 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -41,6 +41,10 @@ impl BufferAllocatorRef { Self(None) } + pub(crate) fn is_statically_allocated(&self) -> bool { + self.0.is_none() + } + /// 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-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 26fc088ff2a..2800fd24214 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -478,23 +478,44 @@ impl BufferMut { .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let old_offset = self.offset; - self.allocation.grow(layout); - let new_offset = self - .allocation - .ptr() - .as_ptr() - .align_offset(physical_alignment.as_usize()); - if new_offset != old_offset { - // SAFETY: grow preserved the initialized elements at old_offset. The new allocation - // has room for the requested elements plus alignment padding, and copy permits overlap. + let new_offset = if self.allocation.allocator().is_statically_allocated() { + let allocation = + Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); + let new_offset = allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + // SAFETY: both allocations have room for the initialized elements and do not overlap. unsafe { - std::ptr::copy( + std::ptr::copy_nonoverlapping( self.allocation.ptr().as_ptr().add(old_offset), - self.allocation.ptr().as_ptr().add(new_offset), + allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); } - } + self.allocation = allocation; + new_offset + } else { + self.allocation.grow(layout); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits + // overlap. + unsafe { + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } + } + new_offset + }; self.offset = new_offset; self.physical_alignment = physical_alignment; } From cbe6a19ab551fdf90883ecc9c664a8d26149edef Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 21:59:21 -0400 Subject: [PATCH 13/29] perf(buffer): double logical growth capacity Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 2800fd24214..1cac9edcfda 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -464,11 +464,13 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.physical_alignment); - let required_size = required_size + let current_size = self.allocation.size() - self.offset; + let logical_size = required_size + .max(current_size.saturating_mul(2)) + .max(physical_alignment.as_usize()); + let allocation_size = logical_size .checked_add(physical_alignment.as_usize()) .vortex_expect("buffer capacity overflow"); - let current_size = self.allocation.size() - self.offset; - let allocation_size = required_size.max(current_size.saturating_mul(2)); let allocation_alignment = if self.allocation.size() == 0 { 1 } else { @@ -1087,6 +1089,19 @@ mod test { assert_eq!(buffer.as_slice(), (0..10_000).collect::>()); } + #[test] + fn growth_seeds_and_doubles_logical_capacity() { + let alignment = Alignment::new(64); + let mut buffer = BufferMut::::empty_aligned(alignment); + + buffer.push(0); + let capacity = buffer.capacity(); + assert!(capacity >= alignment.as_usize()); + + buffer.reserve(capacity); + assert!(buffer.capacity() >= capacity * 2); + } + #[test] fn from_iter() { let buf = BufferMut::from_iter([0, 10, 20, 30]); From a0c8d183936144b98c1e6c643bc35ad4d2265083 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 22:37:46 -0400 Subject: [PATCH 14/29] perf(buffer): exclude alignment slack from growth Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer.rs | 59 ++++++++++++++++++++++++++++++++- vortex-buffer/src/buffer_mut.rs | 27 ++++++++++++--- 2 files changed, 81 insertions(+), 5 deletions(-) diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 026d3bcf492..a453a8ef521 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -35,6 +35,8 @@ pub struct Buffer { pub(crate) length: usize, pub(crate) alignment: Alignment, pub(crate) physical_alignment: Alignment, + // One physical-alignment block is reserved outside the logical capacity. + pub(crate) overallocated: bool, pub(crate) backing: Arc, } @@ -66,6 +68,7 @@ impl Default for Buffer { length: 0, alignment: Alignment::of::(), physical_alignment: Alignment::MAX, + overallocated: false, backing: EMPTY_BACKING.clone(), } } @@ -108,6 +111,7 @@ impl Buffer { length: usize, alignment: Alignment, physical_alignment: Alignment, + overallocated: bool, ) -> Self { // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. let ptr = unsafe { allocation.ptr().add(offset).cast() }; @@ -116,6 +120,7 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing: Arc::new(BufferBacking::Owned(allocation)), } } @@ -133,6 +138,7 @@ impl Buffer { length, alignment, physical_alignment: alignment, + overallocated: false, backing: Arc::new(BufferBacking::External { _owner: owner }), } } @@ -228,6 +234,7 @@ impl Buffer { length: 0, alignment, physical_alignment: Alignment::MAX, + overallocated: false, backing: EMPTY_BACKING.clone(), } } @@ -288,6 +295,7 @@ impl Buffer { length: buffer.length / size_of::(), alignment, physical_alignment: buffer.physical_alignment, + overallocated: buffer.overallocated, backing: buffer.backing, } } @@ -500,6 +508,7 @@ impl Buffer { length: end - begin, alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: Arc::clone(&self.backing), } } @@ -555,6 +564,7 @@ impl Buffer { length: subset.len(), alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: Arc::clone(&self.backing), } } @@ -575,6 +585,7 @@ impl Buffer { length: self.length * size_of::(), alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: self.backing, } } @@ -586,17 +597,25 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing, } = self; match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); + let overallocated = overallocated + && offset + == allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); Ok(BufferMut { allocation, offset, length, alignment, physical_alignment, + overallocated, _marker: Default::default(), }) } @@ -605,6 +624,7 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing: Arc::new(backing), }), Err(backing) => Err(Self { @@ -612,6 +632,7 @@ impl Buffer { length, alignment, physical_alignment, + overallocated, backing, }), } @@ -684,6 +705,7 @@ impl Buffer { length: self.length, alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, backing: self.backing, } } @@ -789,7 +811,14 @@ where if std::mem::needs_drop::() { Self::from_owner(Wrapper(value), alignment) } else { - Self::from_allocation(Allocation::from_vec(value), 0, length, alignment, alignment) + Self::from_allocation( + Allocation::from_vec(value), + 0, + length, + alignment, + alignment, + false, + ) } } } @@ -1055,6 +1084,34 @@ mod test { assert_eq!(buffer.allocation.alignment(), align_of::()); } + #[test] + fn from_u8_vec_preserves_capacity() { + let mut vec = Vec::with_capacity(16); + vec.extend([1u8, 2, 3]); + + let buffer = Buffer::from(vec); + let Ok(buffer) = buffer.try_into_mut() else { + panic!("Vec-backed buffer should be uniquely owned") + }; + assert_eq!(buffer.capacity(), 16); + } + + #[test] + fn sliced_buffer_into_mut_has_safe_capacity() { + let mut original = crate::BufferMut::with_capacity(128); + original.extend(0u32..100); + let original = original.freeze(); + let sliced = original.slice(64..96); + drop(original); + + let Ok(mut sliced) = sliced.try_into_mut() else { + panic!("uniquely owned slice should become mutable") + }; + let capacity = sliced.capacity(); + sliced.push_n(0, capacity - sliced.len()); + assert_eq!(sliced.len(), capacity); + } + #[test] fn from_vec_preserves_drop_glue() { struct DropValue(Arc); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 1cac9edcfda..787587e6a66 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -33,6 +33,8 @@ pub struct BufferMut { pub(crate) length: usize, pub(crate) alignment: Alignment, pub(crate) physical_alignment: Alignment, + // One physical-alignment block is reserved outside the logical capacity. + pub(crate) overallocated: bool, pub(crate) _marker: std::marker::PhantomData, } @@ -134,6 +136,7 @@ impl BufferMut { length: 0, alignment, physical_alignment: actual, + overallocated: true, _marker: Default::default(), } } @@ -222,6 +225,7 @@ impl BufferMut { length: len, alignment, physical_alignment: actual_alignment, + overallocated: true, _marker: Default::default(), } } @@ -390,7 +394,15 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - (self.allocation.size() - self.offset) / size_of::() + if self.allocation.size() == 0 { + return 0; + } + + if !self.overallocated { + return (self.allocation.size() - self.offset) / size_of::(); + } + + (self.allocation.size() - self.physical_alignment.as_usize()) / size_of::() } /// Returns a raw pointer to the buffer's data. @@ -464,7 +476,10 @@ impl BufferMut { .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.physical_alignment); - let current_size = self.allocation.size() - self.offset; + let current_size = self + .capacity() + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); let logical_size = required_size .max(current_size.saturating_mul(2)) .max(physical_alignment.as_usize()); @@ -520,6 +535,7 @@ impl BufferMut { }; self.offset = new_offset; self.physical_alignment = physical_alignment; + self.overallocated = true; } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -669,6 +685,7 @@ impl BufferMut { length: self.length * size_of::(), alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, _marker: Default::default(), } } @@ -681,6 +698,7 @@ impl BufferMut { self.length, self.alignment, self.physical_alignment, + self.overallocated, ) } @@ -745,6 +763,7 @@ impl BufferMut { length: self.length, alignment: self.alignment, physical_alignment: self.physical_alignment, + overallocated: self.overallocated, _marker: std::marker::PhantomData, } } @@ -1096,10 +1115,10 @@ mod test { buffer.push(0); let capacity = buffer.capacity(); - assert!(capacity >= alignment.as_usize()); + assert_eq!(capacity, Alignment::DEFAULT_ALIGNMENT.as_usize()); buffer.reserve(capacity); - assert!(buffer.capacity() >= capacity * 2); + assert_eq!(buffer.capacity(), capacity * 2); } #[test] From b9836166a659fae3ee080f674dae809f5550eadf Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 22:44:59 -0400 Subject: [PATCH 15/29] perf(buffer): store aligned mutable pointer Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer.rs | 2 +- vortex-buffer/src/buffer_mut.rs | 43 +++++++++++++++++++++------------ 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index a453a8ef521..33fc9933ad0 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -611,7 +611,7 @@ impl Buffer { .align_offset(physical_alignment.as_usize()); Ok(BufferMut { allocation, - offset, + ptr, length, alignment, physical_alignment, diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 787587e6a66..a34fc5d374a 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -29,7 +29,7 @@ use crate::trusted_len::TrustedLen; /// A mutable buffer that maintains a runtime-defined alignment through resizing operations. pub struct BufferMut { pub(crate) allocation: Allocation, - pub(crate) offset: usize, + pub(crate) ptr: std::ptr::NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, pub(crate) physical_alignment: Alignment, @@ -38,6 +38,11 @@ pub struct BufferMut { pub(crate) _marker: std::marker::PhantomData, } +// SAFETY: BufferMut uniquely owns its allocation and only exposes T across threads. +unsafe impl Send for BufferMut {} +// SAFETY: shared access to BufferMut only exposes shared access to T. +unsafe impl Sync for BufferMut {} + impl BufferMut { /// Create a new `BufferMut` with the requested alignment and capacity. pub fn with_capacity(capacity: usize) -> Self { @@ -130,9 +135,11 @@ impl BufferMut { }; let allocation = Allocation::allocate(layout, allocator); let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); + // SAFETY: the allocation includes enough padding to reach this aligned pointer. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; Self { allocation, - offset, + ptr, length: 0, alignment, physical_alignment: actual, @@ -219,9 +226,11 @@ impl BufferMut { .ptr() .as_ptr() .align_offset(actual_alignment.as_usize()); + // SAFETY: the allocation includes enough padding to reach this aligned pointer. + let ptr = unsafe { allocation.ptr().add(offset).cast() }; Self { allocation, - offset, + ptr, length: len, alignment, physical_alignment: actual_alignment, @@ -399,7 +408,8 @@ impl BufferMut { } if !self.overallocated { - return (self.allocation.size() - self.offset) / size_of::(); + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + return (self.allocation.size() - offset) / size_of::(); } (self.allocation.size() - self.physical_alignment.as_usize()) / size_of::() @@ -408,21 +418,19 @@ impl BufferMut { /// Returns a raw pointer to the buffer's data. #[inline(always)] pub fn as_ptr(&self) -> *const T { - // SAFETY: offset always remains within the allocation. - unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } + self.ptr.as_ptr() } /// Returns a mutable raw pointer to the buffer's data. #[inline(always)] pub fn as_mut_ptr(&mut self) -> *mut T { - // SAFETY: BufferMut uniquely owns the allocation and offset is in bounds. - unsafe { self.allocation.ptr().as_ptr().add(self.offset).cast() } + self.ptr.as_ptr() } /// Returns a slice over the buffer of elements of type T. #[inline] pub fn as_slice(&self) -> &[T] { - // SAFETY: the allocation is live, offset is in bounds, and construction checks alignment. + // SAFETY: ptr is in the live allocation and construction checks its alignment. unsafe { std::slice::from_raw_parts(self.as_ptr(), self.length) } } @@ -494,7 +502,7 @@ impl BufferMut { let layout = Layout::from_size_align(allocation_size, allocation_alignment) .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); - let old_offset = self.offset; + let old_offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); let new_offset = if self.allocation.allocator().is_statically_allocated() { let allocation = Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); @@ -505,7 +513,7 @@ impl BufferMut { // SAFETY: both allocations have room for the initialized elements and do not overlap. unsafe { std::ptr::copy_nonoverlapping( - self.allocation.ptr().as_ptr().add(old_offset), + self.ptr.cast::().as_ptr(), allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); @@ -533,7 +541,8 @@ impl BufferMut { } new_offset }; - self.offset = new_offset; + // SAFETY: new_offset was computed within the allocation for physical_alignment. + self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; self.physical_alignment = physical_alignment; self.overallocated = true; } @@ -681,7 +690,7 @@ impl BufferMut { pub fn into_byte_buffer(self) -> ByteBufferMut { ByteBufferMut { allocation: self.allocation, - offset: self.offset, + ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, physical_alignment: self.physical_alignment, @@ -692,9 +701,10 @@ impl BufferMut { /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); Buffer::from_allocation( self.allocation, - self.offset, + offset, self.length, self.alignment, self.physical_alignment, @@ -759,7 +769,7 @@ impl BufferMut { BufferMut { allocation: self.allocation, - offset: self.offset, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, physical_alignment: self.physical_alignment, @@ -1003,7 +1013,8 @@ impl Buf for ByteBufferMut { ); } assert!(cnt <= self.length, "advance out of bounds"); - self.offset += cnt; + // SAFETY: cnt is checked against the initialized length above. + self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; } } From 238f6ef23d51c52397118cf251995c2aca698ab0 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Mon, 31 Aug 2026 11:05:24 -0400 Subject: [PATCH 16/29] fix(buffer): tighten allocation ownership paths Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 15 ++- vortex-buffer/src/arrow.rs | 34 +++--- vortex-buffer/src/buffer.rs | 180 ++++++++++++++++++++++++-------- vortex-buffer/src/buffer_mut.rs | 80 +++++++------- 4 files changed, 201 insertions(+), 108 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 415ba4fc4b4..7b405b300b6 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -41,8 +41,8 @@ impl BufferAllocatorRef { Self(None) } - pub(crate) fn is_statically_allocated(&self) -> bool { - self.0.is_none() + pub(crate) fn static_ref() -> &'static Self { + &STATIC_ALLOCATOR } /// Create a mutable buffer with this allocator. @@ -337,7 +337,12 @@ where pub(crate) enum BufferBacking { Owned(Allocation), - External { _owner: Box }, + Bytes(bytes::Bytes), + #[cfg(feature = "arrow")] + Arrow(arrow_buffer::Buffer), + External { + _owner: Box, + }, } impl BufferBacking { @@ -345,7 +350,9 @@ impl BufferBacking { pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { Self::Owned(allocation) => allocation.allocator(), - Self::External { .. } => &STATIC_ALLOCATOR, + Self::Bytes(_) | Self::External { .. } => &STATIC_ALLOCATOR, + #[cfg(feature = "arrow")] + Self::Arrow(_) => &STATIC_ALLOCATOR, } } } diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index f63aaadaff0..c1d2ae23b87 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -3,7 +3,6 @@ use arrow_buffer::ArrowNativeType; use arrow_buffer::OffsetBuffer; -use bytes::Bytes; use vortex_error::vortex_panic; use crate::Alignment; @@ -13,7 +12,7 @@ use crate::ByteBuffer; impl Buffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_scalar_buffer(self) -> arrow_buffer::ScalarBuffer { - let buffer = arrow_buffer::Buffer::from(self.into_inner()); + let buffer = self.into_byte_buffer().into_arrow_buffer(); arrow_buffer::ScalarBuffer::from(buffer) } @@ -25,18 +24,18 @@ impl Buffer { /// alignment is not sufficient for type T. pub fn from_arrow_scalar_buffer(arrow: arrow_buffer::ScalarBuffer) -> Self { let length = arrow.len(); - let bytes = Bytes::from_owner(ArrowWrapper(arrow.into_inner())); + let arrow = arrow.into_inner(); let alignment = Alignment::of::(); - if bytes.as_ptr().align_offset(alignment.as_usize()) != 0 { + if arrow.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment ); } - debug_assert_eq!(length, bytes.len() / size_of::()); - Self::from_bytes_aligned(bytes, alignment) + debug_assert_eq!(length, arrow.len() / size_of::()); + Self::from_arrow_owner(arrow, length, alignment) } /// Converts the buffer zero-copy into a `arrow_buffer::OffsetBuffer`. @@ -51,6 +50,10 @@ impl Buffer { impl ByteBuffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_buffer(self) -> arrow_buffer::Buffer { + if let Some(crate::BufferBacking::Arrow(arrow)) = self.backing.as_deref() { + let offset = self.ptr.addr().get() - arrow.as_ptr().addr(); + return arrow.slice_with_length(offset, self.length); + } arrow_buffer::Buffer::from(self.into_inner()) } @@ -62,26 +65,14 @@ impl ByteBuffer { pub fn from_arrow_buffer(arrow: arrow_buffer::Buffer, alignment: Alignment) -> Self { let length = arrow.len(); - let bytes = Bytes::from_owner(ArrowWrapper(arrow)); - if bytes.as_ptr().align_offset(alignment.as_usize()) != 0 { + if arrow.as_ptr().align_offset(alignment.as_usize()) != 0 { vortex_panic!( "Arrow buffer is not aligned to the requested alignment: {}", alignment ); } - debug_assert_eq!(length, bytes.len()); - Self::from_bytes_aligned(bytes, alignment) - } -} - -/// A wrapper struct to allow `arrow_buffer::Buffer` to implement `AsRef<[u8]>` for -/// `Bytes::from_owner`. -struct ArrowWrapper(arrow_buffer::Buffer); - -impl AsRef<[u8]> for ArrowWrapper { - fn as_ref(&self) -> &[u8] { - self.0.as_slice() + Self::from_arrow_owner(arrow, length, alignment) } } @@ -116,5 +107,8 @@ mod test { let buf = Buffer::from_arrow_buffer(arrow.clone(), Alignment::of::()); assert_eq!(arrow.as_ref(), buf.as_slice(), "Buffer values differ"); assert_eq!(arrow.as_ptr(), buf.as_ptr(), "Conversion not zero-copy"); + + let round_trip = buf.into_arrow_buffer(); + assert_eq!(round_trip.as_ptr(), arrow.as_ptr()); } } diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 33fc9933ad0..dce0cb85e46 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -12,7 +12,6 @@ use std::ops::Deref; use std::ops::RangeBounds; use std::ptr::NonNull; use std::sync::Arc; -use std::sync::LazyLock; use bytes::Buf; use bytes::Bytes; @@ -37,7 +36,7 @@ pub struct Buffer { pub(crate) physical_alignment: Alignment, // One physical-alignment block is reserved outside the logical capacity. pub(crate) overallocated: bool, - pub(crate) backing: Arc, + pub(crate) backing: Option>, } // SAFETY: Buffer is an immutable view over backing memory. Its pointer remains valid while the @@ -46,21 +45,6 @@ unsafe impl Send for Buffer {} // SAFETY: see the Send implementation above. unsafe impl Sync for Buffer {} -/// Zero-length backing for empty buffers, "aligned" to [`Alignment::MAX`] so it satisfies any -/// valid alignment without allocating. A zero-length slice never reads memory, so it may use a -/// dangling pointer as long as it is non-null and aligned. -const EMPTY_BYTES: &[u8] = { - let ptr = std::ptr::without_provenance(Alignment::MAX.as_usize()); - // SAFETY: the pointer is non-null and aligned, and the slice is zero-length. - unsafe { std::slice::from_raw_parts(ptr, 0) } -}; - -static EMPTY_BACKING: LazyLock> = LazyLock::new(|| { - Arc::new(BufferBacking::External { - _owner: Box::new(EMPTY_BYTES), - }) -}); - impl Default for Buffer { fn default() -> Self { Self { @@ -69,7 +53,7 @@ impl Default for Buffer { alignment: Alignment::of::(), physical_alignment: Alignment::MAX, overallocated: false, - backing: EMPTY_BACKING.clone(), + backing: None, } } } @@ -121,7 +105,7 @@ impl Buffer { alignment, physical_alignment, overallocated, - backing: Arc::new(BufferBacking::Owned(allocation)), + backing: Some(Arc::new(BufferBacking::Owned(allocation))), } } @@ -139,7 +123,45 @@ impl Buffer { alignment, physical_alignment: alignment, overallocated: false, - backing: Arc::new(BufferBacking::External { _owner: owner }), + backing: Some(Arc::new(BufferBacking::External { _owner: owner })), + } + } + + fn from_bytes(bytes: Bytes, alignment: Alignment) -> Self { + let length = bytes.len() / size_of::(); + if length == 0 { + return Self::empty_aligned(alignment); + } + let ptr = + NonNull::new(bytes.as_ptr().cast_mut().cast()).vortex_expect("Bytes pointer is null"); + Self { + ptr, + length, + alignment, + physical_alignment: alignment, + overallocated: false, + backing: Some(Arc::new(BufferBacking::Bytes(bytes))), + } + } + + #[cfg(feature = "arrow")] + pub(crate) fn from_arrow_owner( + arrow: arrow_buffer::Buffer, + length: usize, + alignment: Alignment, + ) -> Self { + if length == 0 { + return Self::empty_aligned(alignment); + } + let ptr = NonNull::new(arrow.as_ptr().cast_mut().cast()) + .vortex_expect("Arrow buffer pointer is null"); + Self { + ptr, + length, + alignment, + physical_alignment: alignment, + overallocated: false, + backing: Some(Arc::new(BufferBacking::Arrow(arrow))), } } @@ -219,8 +241,7 @@ impl Buffer { /// Create a new empty `ByteBuffer` with the provided alignment. /// - /// This does not allocate: empty buffers are backed by a zero-length `Bytes` that is - /// aligned to [`Alignment::MAX`]. + /// This does not allocate. Empty buffers use an aligned dangling pointer. pub fn empty_aligned(alignment: Alignment) -> Self { if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( @@ -235,7 +256,7 @@ impl Buffer { alignment, physical_alignment: Alignment::MAX, overallocated: false, - backing: EMPTY_BACKING.clone(), + backing: None, } } @@ -327,7 +348,7 @@ impl Buffer { size_of::() ); } - Self::from_owner(bytes, alignment) + Self::from_bytes(bytes, alignment) } /// Create a buffer with values from the TrustedLen iterator. @@ -346,7 +367,7 @@ impl Buffer { Ok(mut_buf) => mut_buf.map_each_in_place(f), Err(buf) => { let len = buf.len(); - let allocator = buf.backing.allocator().clone(); + let allocator = buf.allocator().clone(); let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf .spare_capacity_mut() @@ -392,7 +413,10 @@ impl Buffer { /// /// External buffers use the static allocator. pub fn allocator(&self) -> &BufferAllocatorRef { - self.backing.allocator() + match self.backing.as_deref() { + Some(backing) => backing.allocator(), + None => BufferAllocatorRef::static_ref(), + } } /// Returns a raw pointer to the buffer's data. @@ -509,7 +533,7 @@ impl Buffer { alignment, physical_alignment: self.physical_alignment, overallocated: self.overallocated, - backing: Arc::clone(&self.backing), + backing: self.backing.clone(), } } @@ -565,17 +589,36 @@ impl Buffer { alignment, physical_alignment: self.physical_alignment, overallocated: self.overallocated, - backing: Arc::clone(&self.backing), + backing: self.backing.clone(), } } /// Returns the underlying bytes without copying. pub fn into_inner(self) -> Bytes { - Bytes::from_owner(BufferBytesOwner { - ptr: self.ptr.cast(), - length: self.length * size_of::(), - backing: self.backing, - }) + if let Some(backing) = self.backing.as_ref() + && let BufferBacking::Bytes(bytes) = backing.as_ref() + { + let offset = self.ptr.cast::().addr().get() - bytes.as_ptr().addr(); + let length = self.length * size_of::(); + if offset == 0 && length == bytes.len() && Arc::strong_count(backing) == 1 { + return match self.backing { + Some(backing) => match Arc::try_unwrap(backing) { + Ok(BufferBacking::Bytes(bytes)) => bytes, + _ => unreachable!(), + }, + None => unreachable!(), + }; + } + return bytes.slice(offset..offset + length); + } + match self.backing { + Some(backing) => Bytes::from_owner(BufferBytesOwner { + ptr: self.ptr.cast(), + length: self.length * size_of::(), + backing, + }), + None => Bytes::new(), + } } /// Return the ByteBuffer for this `Buffer`. @@ -600,6 +643,19 @@ impl Buffer { overallocated, backing, } = self; + let Some(backing) = backing else { + return Ok(BufferMut::empty_aligned(alignment)); + }; + if !matches!(backing.as_ref(), BufferBacking::Owned(_)) { + return Err(Self { + ptr, + length, + alignment, + physical_alignment, + overallocated, + backing: Some(backing), + }); + } match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); @@ -619,21 +675,14 @@ impl Buffer { _marker: Default::default(), }) } - Ok(backing) => Err(Self { - ptr, - length, - alignment, - physical_alignment, - overallocated, - backing: Arc::new(backing), - }), + Ok(_) => unreachable!(), Err(backing) => Err(Self { ptr, length, alignment, physical_alignment, overallocated, - backing, + backing: Some(backing), }), } } @@ -641,7 +690,7 @@ impl Buffer { /// Convert self into `BufferMut`, cloning the data if there are multiple strong references. pub fn into_mut(self) -> BufferMut { self.try_into_mut().unwrap_or_else(|buffer| { - let allocator = buffer.backing.allocator().clone(); + let allocator = buffer.allocator().clone(); BufferMut::::copy_from_aligned_in(&buffer, buffer.alignment, allocator) }) } @@ -664,7 +713,7 @@ impl Buffer { "Buffer is not aligned to requested alignment {alignment}, copying: {bt}" ) } - let allocator = self.backing.allocator().clone(); + let allocator = self.allocator().clone(); BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze() } } @@ -825,7 +874,7 @@ where impl From for ByteBuffer { fn from(bytes: Bytes) -> Self { - Self::from_owner(bytes, Alignment::of::()) + Self::from_bytes(bytes, Alignment::of::()) } } @@ -947,9 +996,11 @@ mod test { use std::sync::atomic::Ordering; use bytes::Buf; + use bytes::Bytes; use crate::Alignment; use crate::Buffer; + use crate::BufferBacking; use crate::ByteBuffer; use crate::buffer; @@ -1084,6 +1135,40 @@ mod test { assert_eq!(buffer.allocation.alignment(), align_of::()); } + #[test] + fn bytes_round_trip_reuses_owner() { + let bytes = Bytes::from_static(&[1, 2, 3, 4]); + let ptr = bytes.as_ptr(); + + let buffer = ByteBuffer::from(bytes); + assert!(matches!( + buffer.backing.as_deref(), + Some(BufferBacking::Bytes(_)) + )); + let bytes = buffer.into_inner(); + + assert_eq!(bytes.as_ptr(), ptr); + assert_eq!(bytes.as_ref(), &[1, 2, 3, 4]); + } + + #[test] + fn external_try_into_mut_preserves_backing() { + let buffer = ByteBuffer::from(Bytes::from_static(&[1, 2, 3, 4])); + let Some(original_backing) = buffer.backing.as_ref() else { + panic!("external buffer has no backing") + }; + let backing = Arc::as_ptr(original_backing); + + let Err(buffer) = buffer.try_into_mut() else { + panic!("external buffer became mutable") + }; + + let Some(new_backing) = buffer.backing.as_ref() else { + panic!("external buffer has no backing") + }; + assert_eq!(Arc::as_ptr(new_backing), backing); + } + #[test] fn from_u8_vec_preserves_capacity() { let mut vec = Vec::with_capacity(16); @@ -1141,6 +1226,11 @@ mod test { assert!(buf.is_aligned(Alignment::MAX)); } + #[test] + fn empty_has_no_backing() { + assert!(Buffer::::empty().backing.is_none()); + } + #[test] fn empty_slice_preserves_alignment() { let buf = Buffer::::zeroed_aligned(8, Alignment::new(64)); diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index a34fc5d374a..0251509bfed 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -503,44 +503,24 @@ impl BufferMut { .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let old_offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); - let new_offset = if self.allocation.allocator().is_statically_allocated() { - let allocation = - Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); - let new_offset = allocation - .ptr() - .as_ptr() - .align_offset(physical_alignment.as_usize()); - // SAFETY: both allocations have room for the initialized elements and do not overlap. + self.allocation.grow(layout); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits + // overlap. unsafe { - std::ptr::copy_nonoverlapping( - self.ptr.cast::().as_ptr(), - allocation.ptr().as_ptr().add(new_offset), + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); } - self.allocation = allocation; - new_offset - } else { - self.allocation.grow(layout); - let new_offset = self - .allocation - .ptr() - .as_ptr() - .align_offset(physical_alignment.as_usize()); - if new_offset != old_offset { - // SAFETY: grow preserved the initialized elements at old_offset. The new allocation - // has room for the requested elements plus alignment padding, and copy permits - // overlap. - unsafe { - std::ptr::copy( - self.allocation.ptr().as_ptr().add(old_offset), - self.allocation.ptr().as_ptr().add(new_offset), - self.length * size_of::(), - ); - } - } - new_offset - }; + } // SAFETY: new_offset was computed within the allocation for physical_alignment. self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; self.physical_alignment = physical_alignment; @@ -737,11 +717,7 @@ impl BufferMut { /// If the data is not aligned, we copy it into a new allocation. pub fn aligned(self, alignment: Alignment) -> Self { if self.as_ptr().align_offset(alignment.as_usize()) == 0 { - Self { - alignment, - physical_alignment: max(self.physical_alignment, alignment), - ..self - } + Self { alignment, ..self } } else { let allocator = self.allocation.allocator().clone(); Self::copy_from_aligned_in(self, alignment, allocator) @@ -1016,6 +992,7 @@ impl Buf for ByteBufferMut { // SAFETY: cnt is checked against the initialized length above. self.ptr = unsafe { self.ptr.add(cnt) }; self.length -= cnt; + self.overallocated = false; } } @@ -1132,6 +1109,31 @@ mod test { assert_eq!(buffer.capacity(), capacity * 2); } + #[test] + fn advance_uses_remaining_allocation_capacity() { + let mut buffer = ByteBufferMut::zeroed_aligned(64, Alignment::new(8)); + + buffer.advance(8); + + assert!(!buffer.overallocated); + let offset = buffer.ptr.addr().get() - buffer.allocation.ptr().addr().get(); + assert_eq!(buffer.capacity(), buffer.allocation.size() - offset); + } + + #[test] + fn raising_logical_alignment_preserves_capacity() { + let buffer = + BufferMut::::with_capacity_preferred_aligned(1, Alignment::of::(), None); + let capacity = buffer.capacity(); + + let mut buffer = buffer.aligned(Alignment::new(2)); + + assert_eq!(buffer.capacity(), capacity); + buffer.extend(0..100); + assert!(Alignment::new(2).is_ptr_aligned(buffer.as_ptr())); + assert_eq!(buffer.as_slice(), (0..100).collect::>()); + } + #[test] fn from_iter() { let buf = BufferMut::from_iter([0, 10, 20, 30]); From 4749a94e728f1b828d6612706de010a433aa54ed Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Mon, 31 Aug 2026 11:53:28 -0400 Subject: [PATCH 17/29] refactor(buffer): remove mutable bytes traits Signed-off-by: Nicholas Gates --- encodings/pco/src/array.rs | 9 +- encodings/sparse/src/lib.rs | 3 +- fuzz/fuzz_targets/file_io.rs | 3 +- .../src/arrays/constant/vtable/mod.rs | 3 +- vortex-buffer/src/buffer_mut.rs | 127 ------------------ 5 files changed, 7 insertions(+), 138 deletions(-) diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 7dc3a371bda..44a6a8e4045 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -47,7 +47,6 @@ use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -568,17 +567,17 @@ impl PcoData { } ); - let mut chunk_meta_buffer = ByteBufferMut::with_capacity(cc.meta_size_hint()); + let mut chunk_meta_buffer = Vec::with_capacity(cc.meta_size_hint()); cc.write_meta(&mut chunk_meta_buffer) .map_err(vortex_err_from_pco)?; - chunk_meta_buffers.push(chunk_meta_buffer.freeze()); + chunk_meta_buffers.push(ByteBuffer::from(chunk_meta_buffer)); let mut page_infos = vec![]; for (page_idx, page_n_values) in cc.n_per_page().into_iter().enumerate() { - let mut page = ByteBufferMut::with_capacity(cc.page_size_hint(page_idx)); + let mut page = Vec::with_capacity(cc.page_size_hint(page_idx)); cc.write_page(page_idx, &mut page) .map_err(vortex_err_from_pco)?; - page_buffers.push(page.freeze()); + page_buffers.push(ByteBuffer::from(page)); page_infos.push(PcoPageInfo { n_values: u32::try_from(page_n_values)?, }); diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index f3a488ea9ba..e18106f4ca4 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -48,7 +48,6 @@ use vortex_array::validity::Validity; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; use vortex_buffer::Buffer; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -217,7 +216,7 @@ impl VTable for Sparse { match idx { 0 => { let fill_value_buffer = - ScalarValue::to_proto_bytes::(array.fill_value.value()).freeze(); + ScalarValue::to_proto_bytes::>(array.fill_value.value()).into(); BufferHandle::new_host(fill_value_buffer) } _ => vortex_panic!("SparseArray buffer index {idx} out of bounds"), diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 38a9016e4b8..6d9c8906fc9 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -19,7 +19,6 @@ use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_btrblocks::BtrBlocksCompressorBuilder; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_file::OpenOptionsSessionExt; @@ -72,7 +71,7 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { ), }; - let mut full_buff = ByteBufferMut::empty(); + let mut full_buff = Vec::new(); let _footer = write_options .blocking(&*RUNTIME) .write(&mut full_buff, array_data.to_array_iterator()) diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 6917d979239..7e950bfe5b5 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -6,7 +6,6 @@ use std::hash::Hash; use std::hash::Hasher; use itertools::Itertools; -use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -107,7 +106,7 @@ impl VTable for Constant { fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { match idx { 0 => BufferHandle::new_host( - ScalarValue::to_proto_bytes::(array.scalar.value()).freeze(), + ScalarValue::to_proto_bytes::>(array.scalar.value()).into(), ), _ => vortex_panic!("ConstantArray buffer index {idx} out of bounds"), } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 0251509bfed..dd867e2e64b 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -7,13 +7,9 @@ use std::any::type_name; use std::cmp::max; use std::fmt::Debug; use std::fmt::Formatter; -use std::io::Write; use std::ops::Deref; use std::ops::DerefMut; -use bytes::Buf; -use bytes::BufMut; -use bytes::buf::UninitSlice; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::vortex_panic; @@ -971,101 +967,10 @@ impl FromIterator for BufferMut { } } -impl Buf for ByteBufferMut { - fn remaining(&self) -> usize { - self.len() - } - - fn chunk(&self) -> &[u8] { - self.as_slice() - } - - fn advance(&mut self, cnt: usize) { - if !self.alignment.is_offset_aligned(cnt) { - vortex_panic!( - "Cannot advance buffer by {} items, resulting alignment is not {}", - cnt, - self.alignment - ); - } - assert!(cnt <= self.length, "advance out of bounds"); - // SAFETY: cnt is checked against the initialized length above. - self.ptr = unsafe { self.ptr.add(cnt) }; - self.length -= cnt; - self.overallocated = false; - } -} - -/// As per the BufMut implementation, we must support internal resizing when -/// asked to extend the buffer. -/// See: -unsafe impl BufMut for ByteBufferMut { - #[inline] - fn remaining_mut(&self) -> usize { - usize::MAX - self.len() - } - - #[inline] - unsafe fn advance_mut(&mut self, cnt: usize) { - if !self.alignment.is_offset_aligned(cnt) { - vortex_panic!( - "Cannot advance buffer by {} items, resulting alignment is not {}", - cnt, - self.alignment - ); - } - self.reserve(cnt); - self.length += cnt; - } - - #[inline] - fn chunk_mut(&mut self) -> &mut UninitSlice { - let spare = self.spare_capacity_mut(); - // SAFETY: spare points to valid uninitialized byte capacity owned by this buffer. - unsafe { UninitSlice::from_raw_parts_mut(spare.as_mut_ptr().cast(), spare.len()) } - } - - fn put(&mut self, mut src: T) - where - Self: Sized, - { - while src.has_remaining() { - let chunk = src.chunk(); - self.extend_from_slice(chunk); - src.advance(chunk.len()); - } - } - - #[inline] - fn put_slice(&mut self, src: &[u8]) { - self.extend_from_slice(src); - } - - #[inline] - fn put_bytes(&mut self, val: u8, cnt: usize) { - self.push_n(val, cnt) - } -} - -impl Write for ByteBufferMut { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.extend_from_slice(buf); - Ok(buf.len()) - } - - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } -} - #[cfg(test)] mod test { - use bytes::Buf; - use bytes::BufMut; - use crate::Alignment; use crate::BufferMut; - use crate::ByteBufferMut; use crate::buffer_mut; #[test] @@ -1109,17 +1014,6 @@ mod test { assert_eq!(buffer.capacity(), capacity * 2); } - #[test] - fn advance_uses_remaining_allocation_capacity() { - let mut buffer = ByteBufferMut::zeroed_aligned(64, Alignment::new(8)); - - buffer.advance(8); - - assert!(!buffer.overallocated); - let offset = buffer.ptr.addr().get() - buffer.allocation.ptr().addr().get(); - assert_eq!(buffer.capacity(), buffer.allocation.size() - offset); - } - #[test] fn raising_logical_alignment_preserves_capacity() { let buffer = @@ -1210,27 +1104,6 @@ mod test { assert_eq!(buf.as_slice(), &[1u32, 2, 3]); } - #[test] - fn bytes_buf() { - let mut buf = ByteBufferMut::copy_from("helloworld".as_bytes()); - assert_eq!(buf.remaining(), 10); - assert_eq!(buf.chunk(), b"helloworld"); - - buf.advance(5); - assert_eq!(buf.remaining(), 5); - assert_eq!(buf.as_slice(), b"world"); - assert_eq!(buf.chunk(), b"world"); - } - - #[test] - fn bytes_buf_mut() { - let mut buf = ByteBufferMut::copy_from("hello".as_bytes()); - assert_eq!(BufMut::remaining_mut(&buf), usize::MAX - 5); - - buf.put_slice(b"world"); - assert_eq!(buf.as_slice(), b"helloworld"); - } - #[test] fn buffer_mut_zeroed() { const LEN: usize = 17; From 31dbdd97cf50e4f13f2aab224e3532ece7469d40 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Mon, 31 Aug 2026 12:02:12 -0400 Subject: [PATCH 18/29] fix(buffer): preserve typed empty alignment Signed-off-by: Nicholas Gates --- vortex-buffer/src/arrow.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index c1d2ae23b87..41e8dbbc733 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -12,6 +12,9 @@ use crate::ByteBuffer; impl Buffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_scalar_buffer(self) -> arrow_buffer::ScalarBuffer { + if self.is_empty() { + return Vec::new().into(); + } let buffer = self.into_byte_buffer().into_arrow_buffer(); arrow_buffer::ScalarBuffer::from(buffer) } @@ -101,6 +104,14 @@ mod test { assert_eq!(scalar.as_ptr(), buf.as_ptr(), "Conversion not zero-copy") } + #[test] + fn empty_into_arrow_scalar_buffer() { + let scalar = Buffer::::empty().into_arrow_scalar_buffer(); + + assert!(scalar.is_empty()); + assert_eq!(scalar.as_ptr().align_offset(align_of::()), 0); + } + #[test] fn from_arrow_buffer() { let arrow = ArrowBuffer::from_vec(vec![0i32, 1, 2]); From aa3814b99a6cd6cfb2696b61536497217b918e5f Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Mon, 31 Aug 2026 12:08:58 -0400 Subject: [PATCH 19/29] fix(buffer): preserve capacity when realigning Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer_mut.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index dd867e2e64b..85c84f57ad0 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -715,8 +715,11 @@ impl BufferMut { if self.as_ptr().align_offset(alignment.as_usize()) == 0 { Self { alignment, ..self } } else { + let capacity = self.capacity(); let allocator = self.allocation.allocator().clone(); - Self::copy_from_aligned_in(self, alignment, allocator) + let mut aligned = Self::with_capacity_aligned_in(capacity, alignment, allocator); + aligned.extend_from_slice(&self); + aligned } } From a1f6c60de62c364ae1e5edece0a349a9b69cc096 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Mon, 31 Aug 2026 12:17:07 -0400 Subject: [PATCH 20/29] fix(arrow): align empty byte views Signed-off-by: Nicholas Gates --- vortex-arrow/src/executor/byte_view.rs | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/vortex-arrow/src/executor/byte_view.rs b/vortex-arrow/src/executor/byte_view.rs index 4406bce8ddc..3026317b511 100644 --- a/vortex-arrow/src/executor/byte_view.rs +++ b/vortex-arrow/src/executor/byte_view.rs @@ -6,12 +6,12 @@ use std::sync::Arc; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::GenericByteViewArray; use arrow_array::types::ByteViewType; -use arrow_buffer::ScalarBuffer; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::Nullability; +use vortex_buffer::Buffer; use vortex_error::VortexResult; use crate::dtype::from_arrow_data_type; @@ -22,8 +22,8 @@ pub fn canonical_varbinview_to_arrow( array: &VarBinViewArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let views = - ScalarBuffer::::from(array.views_handle().as_host().clone().into_arrow_buffer()); + let views = Buffer::::from_byte_buffer(array.views_handle().as_host().clone()) + .into_arrow_scalar_buffer(); let buffers: Vec<_> = array .data_buffers() .iter() @@ -64,3 +64,23 @@ pub(super) fn to_arrow_byte_view( let varbinview = array.execute::(ctx)?; execute_varbinview_to_arrow::(&varbinview, ctx) } + +#[cfg(test)] +mod tests { + use arrow_array::types::StringViewType; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + + use super::*; + + #[test] + fn empty_views_are_aligned() -> VortexResult<()> { + let array = VarBinViewArray::from_iter_str(std::iter::empty::<&str>()); + let mut ctx = array_session().create_execution_ctx(); + + let arrow = canonical_varbinview_to_arrow::(&array, &mut ctx)?; + + assert!(arrow.is_empty()); + Ok(()) + } +} From 1bd48078c83f54b2b2a984854c50bc40fc19cd84 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Mon, 31 Aug 2026 16:05:05 -0400 Subject: [PATCH 21/29] perf(buffer): cache mutable capacity Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer.rs | 8 ++++++++ vortex-buffer/src/buffer_mut.rs | 27 ++++++++++++++++----------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index dce0cb85e46..9c3244c2168 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -665,10 +665,18 @@ impl Buffer { .ptr() .as_ptr() .align_offset(physical_alignment.as_usize()); + let capacity = if allocation.size() == 0 { + 0 + } else if overallocated { + (allocation.size() - physical_alignment.as_usize()) / size_of::() + } else { + (allocation.size() - offset) / size_of::() + }; Ok(BufferMut { allocation, ptr, length, + capacity, alignment, physical_alignment, overallocated, diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 85c84f57ad0..cbb430999a4 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -27,6 +27,7 @@ pub struct BufferMut { pub(crate) allocation: Allocation, pub(crate) ptr: std::ptr::NonNull, pub(crate) length: usize, + pub(crate) capacity: usize, pub(crate) alignment: Alignment, pub(crate) physical_alignment: Alignment, // One physical-alignment block is reserved outside the logical capacity. @@ -137,6 +138,7 @@ impl BufferMut { allocation, ptr, length: 0, + capacity, alignment, physical_alignment: actual, overallocated: true, @@ -228,6 +230,7 @@ impl BufferMut { allocation, ptr, length: len, + capacity: len, alignment, physical_alignment: actual_alignment, overallocated: true, @@ -399,16 +402,7 @@ impl BufferMut { /// Returns the capacity of the buffer. #[inline] pub fn capacity(&self) -> usize { - if self.allocation.size() == 0 { - return 0; - } - - if !self.overallocated { - let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); - return (self.allocation.size() - offset) / size_of::(); - } - - (self.allocation.size() - self.physical_alignment.as_usize()) / size_of::() + self.capacity } /// Returns a raw pointer to the buffer's data. @@ -481,7 +475,7 @@ impl BufferMut { .vortex_expect("buffer capacity overflow"); let physical_alignment = max(self.alignment, self.physical_alignment); let current_size = self - .capacity() + .capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let logical_size = required_size @@ -519,6 +513,7 @@ impl BufferMut { } // SAFETY: new_offset was computed within the allocation for physical_alignment. self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; + self.capacity = logical_size / size_of::(); self.physical_alignment = physical_alignment; self.overallocated = true; } @@ -664,10 +659,19 @@ impl BufferMut { /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + let capacity = if self.allocation.size() == 0 { + 0 + } else if self.overallocated { + self.allocation.size() - self.physical_alignment.as_usize() + } else { + self.allocation.size() - offset + }; ByteBufferMut { allocation: self.allocation, ptr: self.ptr.cast(), length: self.length * size_of::(), + capacity, alignment: self.alignment, physical_alignment: self.physical_alignment, overallocated: self.overallocated, @@ -746,6 +750,7 @@ impl BufferMut { allocation: self.allocation, ptr: self.ptr.cast(), length: self.length, + capacity: self.capacity, alignment: self.alignment, physical_alignment: self.physical_alignment, overallocated: self.overallocated, From ce4fe5c1d0b2a13ea35371677f0838a499fefe53 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Tue, 1 Sep 2026 12:03:36 -0400 Subject: [PATCH 22/29] perf(array): freeze default host buffers directly Signed-off-by: Nicholas Gates --- vortex-array/src/memory.rs | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/memory.rs b/vortex-array/src/memory.rs index b52b3442136..5748f04d525 100644 --- a/vortex-array/src/memory.rs +++ b/vortex-array/src/memory.rs @@ -8,7 +8,6 @@ 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; @@ -244,7 +243,7 @@ impl HostAllocator for DefaultHostAllocator { // SAFETY: We fully initialize this slice before freezing it. unsafe { buffer.set_len(len) }; Ok(WritableHostBuffer::new(Box::new( - DefaultWritableHostBuffer { buffer, alignment }, + DefaultWritableHostBuffer { buffer }, ))) } } @@ -252,18 +251,6 @@ impl HostAllocator for DefaultHostAllocator { #[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 { @@ -272,7 +259,7 @@ impl HostBufferMut for DefaultWritableHostBuffer { } fn alignment(&self) -> Alignment { - self.alignment + self.buffer.alignment() } fn as_mut_slice(&mut self) -> &mut [u8] { @@ -280,9 +267,7 @@ impl HostBufferMut for DefaultWritableHostBuffer { } fn freeze(self: Box) -> ByteBuffer { - let Self { buffer, alignment } = *self; - let bytes = Bytes::from_owner(HostBufferOwner { buffer }); - ByteBuffer::from_bytes_aligned(bytes, alignment) + self.buffer.freeze() } } From 5c4731b1e3f1182dc774526d37e111ca512317b8 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Tue, 1 Sep 2026 12:03:40 -0400 Subject: [PATCH 23/29] perf(buffer): copy live data for static growth Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 4 +++ vortex-buffer/src/buffer_mut.rs | 62 +++++++++++++++++++++++++-------- 2 files changed, 52 insertions(+), 14 deletions(-) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index 7b405b300b6..d2ec67b862f 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -45,6 +45,10 @@ impl BufferAllocatorRef { &STATIC_ALLOCATOR } + pub(crate) fn is_statically_allocated(&self) -> bool { + self.0.is_none() + } + /// 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-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index cbb430999a4..84520e1fb1a 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -493,24 +493,44 @@ impl BufferMut { .unwrap_or_else(|_| vortex_panic!("buffer capacity exceeds maximum allocation size")); let old_offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); - self.allocation.grow(layout); - let new_offset = self - .allocation - .ptr() - .as_ptr() - .align_offset(physical_alignment.as_usize()); - if new_offset != old_offset { - // SAFETY: grow preserved the initialized elements at old_offset. The new allocation - // has room for the requested elements plus alignment padding, and copy permits - // overlap. + let new_offset = if self.allocation.allocator().is_statically_allocated() { + let allocation = + Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); + let new_offset = allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + // SAFETY: both allocations have room for the initialized elements and do not overlap. unsafe { - std::ptr::copy( - self.allocation.ptr().as_ptr().add(old_offset), - self.allocation.ptr().as_ptr().add(new_offset), + std::ptr::copy_nonoverlapping( + self.ptr.cast::().as_ptr(), + allocation.ptr().as_ptr().add(new_offset), self.length * size_of::(), ); } - } + self.allocation = allocation; + new_offset + } else { + self.allocation.grow(layout); + let new_offset = self + .allocation + .ptr() + .as_ptr() + .align_offset(physical_alignment.as_usize()); + if new_offset != old_offset { + // SAFETY: grow preserved the initialized elements at old_offset. The new allocation + // has room for the requested elements plus alignment padding, and copy permits + // overlap. + unsafe { + std::ptr::copy( + self.allocation.ptr().as_ptr().add(old_offset), + self.allocation.ptr().as_ptr().add(new_offset), + self.length * size_of::(), + ); + } + } + new_offset + }; // SAFETY: new_offset was computed within the allocation for physical_alignment. self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; self.capacity = logical_size / size_of::(); @@ -1022,6 +1042,20 @@ mod test { assert_eq!(buffer.capacity(), capacity * 2); } + #[test] + fn static_growth_copies_live_data() { + let mut buffer = BufferMut::::with_capacity(1); + let capacity = buffer.capacity(); + buffer.extend(std::iter::repeat_n(7, capacity)); + let old_ptr = buffer.as_ptr(); + + buffer.push(u32::MAX); + + assert_ne!(buffer.as_ptr(), old_ptr); + assert_eq!(&buffer[..capacity], vec![7; capacity]); + assert_eq!(buffer[capacity], u32::MAX); + } + #[test] fn raising_logical_alignment_preserves_capacity() { let buffer = From a45d41c1b59908eade0b7a08578b01b554c48158 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Tue, 1 Sep 2026 14:16:32 -0400 Subject: [PATCH 24/29] perf(array): inline varbin view compaction Signed-off-by: Nicholas Gates --- vortex-array/src/builders/varbinview.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 3a32c2e9e6b..1c33fd14b2d 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -800,7 +800,7 @@ impl ArrayBuilder for VarBinViewBuilder { } impl VarBinViewBuilder { - #[inline] + #[inline(always)] fn push_view( &mut self, view: BinaryView, From 636059087b07d95486e7018c087e6c289e285952 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 2 Sep 2026 09:58:00 -0400 Subject: [PATCH 25/29] perf(array): avoid struct scalar vec realloc Signed-off-by: Nicholas Gates --- .../src/arrays/struct_/vtable/operations.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/arrays/struct_/vtable/operations.rs b/vortex-array/src/arrays/struct_/vtable/operations.rs index d59bf7eeca3..b491e231520 100644 --- a/vortex-array/src/arrays/struct_/vtable/operations.rs +++ b/vortex-array/src/arrays/struct_/vtable/operations.rs @@ -9,6 +9,7 @@ use crate::array::OperationsVTable; use crate::arrays::Struct; use crate::arrays::struct_::StructArrayExt; use crate::scalar::Scalar; +use crate::scalar::ScalarValue; impl OperationsVTable for Struct { fn scalar_at( @@ -16,12 +17,17 @@ impl OperationsVTable for Struct { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let field_scalars: VortexResult> = array + let field_values = array .iter_unmasked_fields() - .map(|field| field.execute_scalar(index, ctx)) - .collect(); + .map(|field| field.execute_scalar(index, ctx).map(Scalar::into_value)) + .collect::>>()?; // SAFETY: The vtable guarantees index is in-bounds and non-null before this is called. - // Each field's scalar_at returns a scalar with the field's own dtype. - Ok(unsafe { Scalar::struct_unchecked(array.dtype().clone(), field_scalars?) }) + // Each field's scalar_at returns a value with the field's own dtype. + Ok(unsafe { + Scalar::new_unchecked( + array.dtype().clone(), + Some(ScalarValue::Tuple(field_values)), + ) + }) } } From 43c10ad2fa13bc9def9659568e2491d2f233ca33 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 2 Sep 2026 13:01:53 -0400 Subject: [PATCH 26/29] fix(buffer): allow intentional forced inlining Signed-off-by: Nicholas Gates --- vortex-buffer/src/allocation.rs | 5 +++++ vortex-buffer/src/buffer.rs | 1 + vortex-buffer/src/buffer_mut.rs | 2 ++ 3 files changed, 8 insertions(+) diff --git a/vortex-buffer/src/allocation.rs b/vortex-buffer/src/allocation.rs index d2ec67b862f..4662c4cfd75 100644 --- a/vortex-buffer/src/allocation.rs +++ b/vortex-buffer/src/allocation.rs @@ -276,21 +276,25 @@ impl Allocation { } } + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn ptr(&self) -> NonNull { self.ptr } + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn size(&self) -> usize { self.layout.size() } + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn alignment(&self) -> usize { self.layout.align() } + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { &self.allocator @@ -350,6 +354,7 @@ pub(crate) enum BufferBacking { } impl BufferBacking { + #[allow(clippy::inline_always)] #[inline(always)] pub(crate) fn allocator(&self) -> &BufferAllocatorRef { match self { diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 9c3244c2168..efd73f903de 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -420,6 +420,7 @@ impl Buffer { } /// Returns a raw pointer to the buffer's data. + #[allow(clippy::inline_always)] #[inline(always)] pub fn as_ptr(&self) -> *const T { self.ptr.as_ptr() diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 84520e1fb1a..4b276a1ee6c 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -406,12 +406,14 @@ impl BufferMut { } /// Returns a raw pointer to the buffer's data. + #[allow(clippy::inline_always)] #[inline(always)] pub fn as_ptr(&self) -> *const T { self.ptr.as_ptr() } /// Returns a mutable raw pointer to the buffer's data. + #[allow(clippy::inline_always)] #[inline(always)] pub fn as_mut_ptr(&mut self) -> *mut T { self.ptr.as_ptr() From 057234eea8f70d34c1aa17385e3555f5ad5452f9 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 2 Sep 2026 13:26:56 -0400 Subject: [PATCH 27/29] fix(array): allow intentional forced inlining Signed-off-by: Nicholas Gates --- vortex-array/src/builders/varbinview.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 1c33fd14b2d..cde20d43019 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -800,6 +800,7 @@ impl ArrayBuilder for VarBinViewBuilder { } impl VarBinViewBuilder { + #[allow(clippy::inline_always)] #[inline(always)] fn push_view( &mut self, From f1a0055304dca6df54f370ebc88d8e4a11fa29c1 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 2 Sep 2026 15:49:31 -0400 Subject: [PATCH 28/29] bench(buffer): compare allocation ownership Signed-off-by: Nicholas Gates --- vortex-buffer/Cargo.toml | 4 ++ vortex-buffer/benches/allocation.rs | 106 ++++++++++++++++++++++++++++ vortex-buffer/src/arrow.rs | 2 +- vortex-buffer/src/buffer.rs | 4 +- 4 files changed, 113 insertions(+), 3 deletions(-) create mode 100644 vortex-buffer/benches/allocation.rs diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index c77d07f253d..02cf4ce7782 100644 --- a/vortex-buffer/Cargo.toml +++ b/vortex-buffer/Cargo.toml @@ -58,3 +58,7 @@ harness = false [[bench]] name = "collect_bool" harness = false + +[[bench]] +name = "allocation" +harness = false diff --git a/vortex-buffer/benches/allocation.rs b/vortex-buffer/benches/allocation.rs new file mode 100644 index 00000000000..6cd923b5dfb --- /dev/null +++ b/vortex-buffer/benches/allocation.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use allocator_api2::alloc::Global; +use arrow_buffer::MutableBuffer; +use bytes::BytesMut; +use divan::Bencher; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; + +const SIZES: &[usize] = &[0, 64, 256, 1024, 16_384, 65_536]; + +fn main() { + divan::main(); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_vortex(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BufferMut::::with_capacity(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_vortex_custom(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| BufferAllocatorRef::new(Global)) + .bench_refs(|allocator| drop(allocator.with_capacity::(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) { + bencher.bench(|| { + drop(BufferMut::::with_capacity_preferred_aligned( + size, + Alignment::of::(), + None, + )) + }); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_bytes(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BytesMut::with_capacity(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_drop_arrow(bencher: Bencher, size: usize) { + bencher.bench(|| drop(MutableBuffer::with_capacity(size))); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_vortex(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BufferMut::::with_capacity(size).freeze())); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_vortex_custom(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| BufferAllocatorRef::new(Global)) + .bench_refs(|allocator| drop(allocator.with_capacity::(size).freeze())); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) { + bencher.bench(|| { + drop( + BufferMut::::with_capacity_preferred_aligned(size, Alignment::of::(), None) + .freeze(), + ) + }); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_bytes(bencher: Bencher, size: usize) { + bencher.bench(|| drop(BytesMut::with_capacity(size).freeze())); +} + +#[divan::bench(args = SIZES)] +fn allocate_freeze_drop_arrow(bencher: Bencher, size: usize) { + bencher.bench(|| { + let buffer: arrow_buffer::Buffer = MutableBuffer::with_capacity(size).into(); + drop(buffer) + }); +} + +#[divan::bench(args = SIZES)] +fn from_vec_drop_vortex(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| vec![0u8; size]) + .bench_values(|values| drop(Buffer::from(values))); +} + +#[divan::bench(args = SIZES)] +fn from_vec_drop_bytes(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| vec![0u8; size]) + .bench_values(|values| drop(bytes::Bytes::from(values))); +} + +#[divan::bench(args = SIZES)] +fn from_vec_drop_arrow(bencher: Bencher, size: usize) { + bencher + .with_inputs(|| vec![0u8; size]) + .bench_values(|values| drop(arrow_buffer::Buffer::from_vec(values))); +} diff --git a/vortex-buffer/src/arrow.rs b/vortex-buffer/src/arrow.rs index 41e8dbbc733..424aa1878c0 100644 --- a/vortex-buffer/src/arrow.rs +++ b/vortex-buffer/src/arrow.rs @@ -57,7 +57,7 @@ impl ByteBuffer { let offset = self.ptr.addr().get() - arrow.as_ptr().addr(); return arrow.slice_with_length(offset, self.length); } - arrow_buffer::Buffer::from(self.into_inner()) + arrow_buffer::Buffer::from(self.into_bytes()) } /// Convert an Arrow scalar buffer into a Vortex scalar buffer. diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index efd73f903de..2e9edfddd41 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -595,7 +595,7 @@ impl Buffer { } /// Returns the underlying bytes without copying. - pub fn into_inner(self) -> Bytes { + pub fn into_bytes(self) -> Bytes { if let Some(backing) = self.backing.as_ref() && let BufferBacking::Bytes(bytes) = backing.as_ref() { @@ -1154,7 +1154,7 @@ mod test { buffer.backing.as_deref(), Some(BufferBacking::Bytes(_)) )); - let bytes = buffer.into_inner(); + let bytes = buffer.into_bytes(); assert_eq!(bytes.as_ptr(), ptr); assert_eq!(bytes.as_ref(), &[1, 2, 3, 4]); From 78491fbda926fac67231634e4b361d1226382ded Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Wed, 2 Sep 2026 15:59:56 -0400 Subject: [PATCH 29/29] refactor(buffer): remove physical alignment state Signed-off-by: Nicholas Gates --- vortex-buffer/src/buffer.rs | 52 +--------------------------- vortex-buffer/src/buffer_mut.rs | 61 ++++++++++++--------------------- 2 files changed, 23 insertions(+), 90 deletions(-) diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 2e9edfddd41..3b48889fa8d 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -33,9 +33,6 @@ pub struct Buffer { pub(crate) ptr: NonNull, pub(crate) length: usize, pub(crate) alignment: Alignment, - pub(crate) physical_alignment: Alignment, - // One physical-alignment block is reserved outside the logical capacity. - pub(crate) overallocated: bool, pub(crate) backing: Option>, } @@ -51,8 +48,6 @@ impl Default for Buffer { ptr: empty_ptr(), length: 0, alignment: Alignment::of::(), - physical_alignment: Alignment::MAX, - overallocated: false, backing: None, } } @@ -94,8 +89,6 @@ impl Buffer { offset: usize, length: usize, alignment: Alignment, - physical_alignment: Alignment, - overallocated: bool, ) -> Self { // SAFETY: BufferMut keeps offset within allocation, including for empty buffers. let ptr = unsafe { allocation.ptr().add(offset).cast() }; @@ -103,8 +96,6 @@ impl Buffer { ptr, length, alignment, - physical_alignment, - overallocated, backing: Some(Arc::new(BufferBacking::Owned(allocation))), } } @@ -121,8 +112,6 @@ impl Buffer { ptr, length, alignment, - physical_alignment: alignment, - overallocated: false, backing: Some(Arc::new(BufferBacking::External { _owner: owner })), } } @@ -138,8 +127,6 @@ impl Buffer { ptr, length, alignment, - physical_alignment: alignment, - overallocated: false, backing: Some(Arc::new(BufferBacking::Bytes(bytes))), } } @@ -159,8 +146,6 @@ impl Buffer { ptr, length, alignment, - physical_alignment: alignment, - overallocated: false, backing: Some(Arc::new(BufferBacking::Arrow(arrow))), } } @@ -254,8 +239,6 @@ impl Buffer { ptr: empty_ptr(), length: 0, alignment, - physical_alignment: Alignment::MAX, - overallocated: false, backing: None, } } @@ -315,8 +298,6 @@ impl Buffer { ptr: buffer.ptr.cast(), length: buffer.length / size_of::(), alignment, - physical_alignment: buffer.physical_alignment, - overallocated: buffer.overallocated, backing: buffer.backing, } } @@ -532,8 +513,6 @@ impl Buffer { ptr: unsafe { self.ptr.add(begin) }, length: end - begin, alignment, - physical_alignment: self.physical_alignment, - overallocated: self.overallocated, backing: self.backing.clone(), } } @@ -588,8 +567,6 @@ impl Buffer { ptr: NonNull::new(subset.as_ptr().cast_mut()).vortex_expect("slice pointer is null"), length: subset.len(), alignment, - physical_alignment: self.physical_alignment, - overallocated: self.overallocated, backing: self.backing.clone(), } } @@ -628,8 +605,6 @@ impl Buffer { ptr: self.ptr.cast(), length: self.length * size_of::(), alignment: self.alignment, - physical_alignment: self.physical_alignment, - overallocated: self.overallocated, backing: self.backing, } } @@ -640,8 +615,6 @@ impl Buffer { ptr, length, alignment, - physical_alignment, - overallocated, backing, } = self; let Some(backing) = backing else { @@ -652,24 +625,14 @@ impl Buffer { ptr, length, alignment, - physical_alignment, - overallocated, backing: Some(backing), }); } match Arc::try_unwrap(backing) { Ok(BufferBacking::Owned(allocation)) => { let offset = ptr.addr().get() - allocation.ptr().addr().get(); - let overallocated = overallocated - && offset - == allocation - .ptr() - .as_ptr() - .align_offset(physical_alignment.as_usize()); let capacity = if allocation.size() == 0 { 0 - } else if overallocated { - (allocation.size() - physical_alignment.as_usize()) / size_of::() } else { (allocation.size() - offset) / size_of::() }; @@ -679,8 +642,6 @@ impl Buffer { length, capacity, alignment, - physical_alignment, - overallocated, _marker: Default::default(), }) } @@ -689,8 +650,6 @@ impl Buffer { ptr, length, alignment, - physical_alignment, - overallocated, backing: Some(backing), }), } @@ -762,8 +721,6 @@ impl Buffer { ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, - physical_alignment: self.physical_alignment, - overallocated: self.overallocated, backing: self.backing, } } @@ -869,14 +826,7 @@ where if std::mem::needs_drop::() { Self::from_owner(Wrapper(value), alignment) } else { - Self::from_allocation( - Allocation::from_vec(value), - 0, - length, - alignment, - alignment, - false, - ) + Self::from_allocation(Allocation::from_vec(value), 0, length, alignment) } } } diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 4b276a1ee6c..b94ca59bc40 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -29,9 +29,6 @@ pub struct BufferMut { pub(crate) length: usize, pub(crate) capacity: usize, pub(crate) alignment: Alignment, - pub(crate) physical_alignment: Alignment, - // One physical-alignment block is reserved outside the logical capacity. - pub(crate) overallocated: bool, pub(crate) _marker: std::marker::PhantomData, } @@ -134,14 +131,17 @@ impl BufferMut { let offset = allocation.ptr().as_ptr().align_offset(actual.as_usize()); // SAFETY: the allocation includes enough padding to reach this aligned pointer. let ptr = unsafe { allocation.ptr().add(offset).cast() }; + let capacity = if size_of::() == 0 { + capacity + } else { + (allocation.size() - offset) / size_of::() + }; Self { allocation, ptr, length: 0, capacity, alignment, - physical_alignment: actual, - overallocated: true, _marker: Default::default(), } } @@ -226,14 +226,17 @@ impl BufferMut { .align_offset(actual_alignment.as_usize()); // SAFETY: the allocation includes enough padding to reach this aligned pointer. let ptr = unsafe { allocation.ptr().add(offset).cast() }; + let capacity = if size_of::() == 0 { + len + } else { + (allocation.size() - offset) / size_of::() + }; Self { allocation, ptr, length: len, - capacity: len, + capacity, alignment, - physical_alignment: actual_alignment, - overallocated: true, _marker: Default::default(), } } @@ -475,16 +478,16 @@ impl BufferMut { let required_size = required .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); - let physical_alignment = max(self.alignment, self.physical_alignment); + let alignment = self.alignment; let current_size = self .capacity .checked_mul(size_of::()) .vortex_expect("buffer capacity overflow"); let logical_size = required_size .max(current_size.saturating_mul(2)) - .max(physical_alignment.as_usize()); + .max(Alignment::DEFAULT_ALIGNMENT.as_usize()); let allocation_size = logical_size - .checked_add(physical_alignment.as_usize()) + .checked_add(alignment.as_usize()) .vortex_expect("buffer capacity overflow"); let allocation_alignment = if self.allocation.size() == 0 { 1 @@ -498,10 +501,7 @@ impl BufferMut { let new_offset = if self.allocation.allocator().is_statically_allocated() { let allocation = Allocation::allocate(layout, BufferAllocatorRef::statically_allocated()); - let new_offset = allocation - .ptr() - .as_ptr() - .align_offset(physical_alignment.as_usize()); + let new_offset = allocation.ptr().as_ptr().align_offset(alignment.as_usize()); // SAFETY: both allocations have room for the initialized elements and do not overlap. unsafe { std::ptr::copy_nonoverlapping( @@ -518,7 +518,7 @@ impl BufferMut { .allocation .ptr() .as_ptr() - .align_offset(physical_alignment.as_usize()); + .align_offset(alignment.as_usize()); if new_offset != old_offset { // SAFETY: grow preserved the initialized elements at old_offset. The new allocation // has room for the requested elements plus alignment padding, and copy permits @@ -533,11 +533,9 @@ impl BufferMut { } new_offset }; - // SAFETY: new_offset was computed within the allocation for physical_alignment. + // SAFETY: new_offset was computed within the allocation for alignment. self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; self.capacity = logical_size / size_of::(); - self.physical_alignment = physical_alignment; - self.overallocated = true; } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -681,22 +679,16 @@ impl BufferMut { /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { - let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); - let capacity = if self.allocation.size() == 0 { - 0 - } else if self.overallocated { - self.allocation.size() - self.physical_alignment.as_usize() - } else { - self.allocation.size() - offset - }; + let capacity = self + .capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); ByteBufferMut { allocation: self.allocation, ptr: self.ptr.cast(), length: self.length * size_of::(), capacity, alignment: self.alignment, - physical_alignment: self.physical_alignment, - overallocated: self.overallocated, _marker: Default::default(), } } @@ -704,14 +696,7 @@ impl BufferMut { /// Freeze the `BufferMut` into a `Buffer`. pub fn freeze(self) -> Buffer { let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); - Buffer::from_allocation( - self.allocation, - offset, - self.length, - self.alignment, - self.physical_alignment, - self.overallocated, - ) + Buffer::from_allocation(self.allocation, offset, self.length, self.alignment) } /// Map each element of the buffer with a closure. @@ -774,8 +759,6 @@ impl BufferMut { length: self.length, capacity: self.capacity, alignment: self.alignment, - physical_alignment: self.physical_alignment, - overallocated: self.overallocated, _marker: std::marker::PhantomData, } }