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..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" 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-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)), + ) + }) } } diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 3a32c2e9e6b..cde20d43019 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -800,7 +800,8 @@ impl ArrayBuilder for VarBinViewBuilder { } impl VarBinViewBuilder { - #[inline] + #[allow(clippy::inline_always)] + #[inline(always)] fn push_view( &mut self, view: BinaryView, 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() } } 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(()) + } +} diff --git a/vortex-buffer/Cargo.toml b/vortex-buffer/Cargo.toml index 705c992f87d..02cf4ce7782 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 } @@ -57,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/allocation.rs b/vortex-buffer/src/allocation.rs new file mode 100644 index 00000000000..4662c4cfd75 --- /dev/null +++ b/vortex-buffer/src/allocation.rs @@ -0,0 +1,481 @@ +// 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::mem::ManuallyDrop; +use std::ptr::NonNull; +use std::sync::Arc; + +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; + +/// 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 {} + +/// A shared reference to a buffer allocator. +#[derive(Clone)] +pub struct BufferAllocatorRef(Option>); + +impl BufferAllocatorRef { + /// Wrap an allocator in a shared reference. + pub fn new(allocator: impl BufferAllocator) -> Self { + Self(Some(Arc::new(allocator))) + } + + /// Return a shared reference to the static allocator. + pub fn statically_allocated() -> Self { + Self(None) + } + + pub(crate) fn static_ref() -> &'static Self { + &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()) + } + + /// 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 { + 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> { + match &self.0 { + Some(allocator) => allocator.allocate(layout), + None => Global.allocate(layout), + } + } + + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + 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. + match &self.0 { + Some(allocator) => unsafe { allocator.deallocate(ptr, layout) }, + None => 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. + 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( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + 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( + &self, + ptr: NonNull, + old_layout: Layout, + new_layout: Layout, + ) -> Result, AllocError> { + // SAFETY: the caller upholds the Allocator contract. + match &self.0 { + Some(allocator) => unsafe { allocator.shrink(ptr, old_layout, new_layout) }, + None => unsafe { Global.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: BufferAllocatorRef = BufferAllocatorRef(None); + +pub(crate) struct Allocation { + ptr: NonNull, + layout: Layout, + 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) + } + + 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, + 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, + allocator, + } + } + + #[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 + } + + 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 { + 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_ptr(&self) -> *const u8; + + fn len(&self) -> usize; +} + +impl BufferOwner for T +where + T: AsRef<[u8]> + Send + Sync + 'static, +{ + fn as_ptr(&self) -> *const u8 { + self.as_ref().as_ptr() + } + + fn len(&self) -> usize { + self.as_ref().len() + } +} + +pub(crate) enum BufferBacking { + Owned(Allocation), + Bytes(bytes::Bytes), + #[cfg(feature = "arrow")] + Arrow(arrow_buffer::Buffer), + External { + _owner: Box, + }, +} + +impl BufferBacking { + #[allow(clippy::inline_always)] + #[inline(always)] + pub(crate) fn allocator(&self) -> &BufferAllocatorRef { + match self { + Self::Owned(allocation) => allocation.allocator(), + Self::Bytes(_) | Self::External { .. } => &STATIC_ALLOCATOR, + #[cfg(feature = "arrow")] + Self::Arrow(_) => &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, + grows: 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) } + } + + 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] + 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::of::().as_usize() + ); + drop(buffer); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 0); + drop(view); + assert_eq!(state.deallocations.load(Ordering::Relaxed), 1); + } + + #[test] + 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); + 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), 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); + } + + #[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/arrow.rs b/vortex-buffer/src/arrow.rs index aa96cab2e84..424aa1878c0 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,10 @@ 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()); + if self.is_empty() { + return Vec::new().into(); + } + let buffer = self.into_byte_buffer().into_arrow_buffer(); arrow_buffer::ScalarBuffer::from(buffer) } @@ -25,22 +27,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 ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + 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`. @@ -55,7 +53,11 @@ impl Buffer { impl ByteBuffer { /// Converts the buffer zero-copy into a `arrow_buffer::Buffer`. pub fn into_arrow_buffer(self) -> arrow_buffer::Buffer { - arrow_buffer::Buffer::from(self.into_inner()) + 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_bytes()) } /// Convert an Arrow scalar buffer into a Vortex scalar buffer. @@ -66,30 +68,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 ); } - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } - } -} - -/// 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) } } @@ -118,11 +104,22 @@ 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]); 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/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 8c78b9f8a1c..3b48889fa8d 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -8,9 +8,10 @@ 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 bytes::Buf; use bytes::Bytes; @@ -18,6 +19,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 +30,126 @@ 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: Option>, } -/// 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] = { - 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) } -}; +// 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 {} 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: None, } } } -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: Some(Arc::new(BufferBacking::Owned(allocation))), + } + } + + fn from_owner(owner: impl crate::BufferOwner, alignment: Alignment) -> Self { + let owner: Box = Box::new(owner); + let length = owner.len() / size_of::(); + let ptr = if length == 0 { + empty_ptr() + } else { + NonNull::new(owner.as_ptr().cast_mut().cast()).vortex_expect("owner pointer is null") + }; + Self { + ptr, + length, + alignment, + 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, + 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, + backing: Some(Arc::new(BufferBacking::Arrow(arrow))), + } + } + /// 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 +160,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 +192,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 @@ -149,8 +226,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!( @@ -160,10 +236,10 @@ impl Buffer { ); } Self { - bytes: Bytes::from_static(EMPTY_BACKING), + ptr: empty_ptr(), length: 0, alignment, - _marker: PhantomData, + backing: None, } } @@ -175,6 +251,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 +277,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 +329,7 @@ impl Buffer { size_of::() ); } - let length = bytes.len() / size_of::(); - Self { - bytes, - length, - alignment, - _marker: Default::default(), - } + Self::from_bytes(bytes, alignment) } /// Create a buffer with values from the TrustedLen iterator. @@ -248,7 +348,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.allocator().clone(); + let mut out_buf = BufferMut::with_capacity_in(len, allocator); out_buf .spare_capacity_mut() .iter_mut() @@ -265,7 +366,6 @@ impl Buffer { /// Clear the buffer, preserving existing capacity. pub fn clear(&mut self) { - self.bytes.clear(); self.length = 0; } @@ -290,19 +390,39 @@ impl Buffer { self.alignment } + /// Returns the allocator to use for derived buffers. + /// + /// External buffers use the static allocator. + pub fn allocator(&self) -> &BufferAllocatorRef { + match self.backing.as_deref() { + Some(backing) => backing.allocator(), + None => BufferAllocatorRef::static_ref(), + } + } + + /// 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 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 +498,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 +509,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: self.backing.clone(), } } @@ -434,74 +553,119 @@ 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: self.backing.clone(), } } - /// 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. - 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 + /// Returns the underlying bytes without copying. + pub fn into_bytes(self) -> Bytes { + 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`. 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; + let Some(backing) = backing else { + return Ok(BufferMut::empty_aligned(alignment)); + }; + if !matches!(backing.as_ref(), BufferBacking::Owned(_)) { + return Err(Self { + ptr, + length, + alignment, + backing: Some(backing), + }); + } + match Arc::try_unwrap(backing) { + Ok(BufferBacking::Owned(allocation)) => { + let offset = ptr.addr().get() - allocation.ptr().addr().get(); + let capacity = if allocation.size() == 0 { + 0 + } else { + (allocation.size() - offset) / size_of::() + }; + Ok(BufferMut { + allocation, + ptr, + length, + capacity, + alignment, + _marker: Default::default(), + }) + } + Ok(_) => unreachable!(), + Err(backing) => Err(Self { + ptr, + length, + alignment, + backing: Some(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.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 +681,8 @@ impl Buffer { "Buffer is not aligned to requested alignment {alignment}, copying: {bt}" ) } - Self::copy_from_aligned(self, alignment) + let allocator = self.allocator().clone(); + BufferMut::copy_from_aligned_in(self, alignment, allocator).freeze() } } @@ -553,10 +718,10 @@ impl Buffer { ); Buffer { - bytes: self.bytes, + ptr: self.ptr.cast(), length: self.length, alignment: self.alignment, - _marker: PhantomData, + backing: self.backing, } } } @@ -637,48 +802,38 @@ 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::() } } 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, + 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) } } } impl From for ByteBuffer { fn from(bytes: Bytes) -> Self { - let length = bytes.len(); - Self { - bytes, - length, - alignment: Alignment::of::(), - _marker: Default::default(), - } + Self::from_bytes(bytes, Alignment::of::()) } } @@ -702,11 +857,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. @@ -769,10 +949,17 @@ 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 bytes::Bytes; use crate::Alignment; use crate::Buffer; + use crate::BufferBacking; use crate::ByteBuffer; use crate::buffer; @@ -886,6 +1073,110 @@ 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 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_bytes(); + + 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); + 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); + + 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. @@ -894,6 +1185,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 4cd51bf2546..b94ca59bc40 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -2,41 +2,50 @@ // 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; use std::fmt::Formatter; -use std::io::Write; use std::ops::Deref; 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) ptr: std::ptr::NonNull, pub(crate) length: usize, + pub(crate) capacity: usize, pub(crate) alignment: Alignment, 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 { - 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 +55,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 +84,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 +113,34 @@ 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 = 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()); + // 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 { - bytes, + allocation, + ptr, length: 0, + capacity, alignment, _marker: Default::default(), } @@ -88,7 +148,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 +163,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 +188,54 @@ 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 = 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() + .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() }; + let capacity = if size_of::() == 0 { + len + } else { + (allocation.size() - offset) / size_of::() + }; Self { - bytes, - length: actual_len, + allocation, + ptr, + length: len, + capacity, alignment, _marker: Default::default(), } @@ -136,7 +253,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 +269,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 +282,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 +316,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 +345,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 +383,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 +405,40 @@ 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. + #[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() } /// 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: ptr is in the live allocation and construction checks its 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 +460,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 +469,73 @@ 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(); - // 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 required = self + .length + .checked_add(additional) + .vortex_expect("buffer capacity overflow"); + let required_size = required + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); + 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(Alignment::DEFAULT_ALIGNMENT.as_usize()); + let allocation_size = logical_size + .checked_add(alignment.as_usize()) + .vortex_expect("buffer capacity overflow"); + 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.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(alignment.as_usize()); + // SAFETY: both allocations have room for the initialized elements and do not overlap. + unsafe { + 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(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 alignment. + self.ptr = unsafe { self.allocation.ptr().add(new_offset).cast() }; + self.capacity = logical_size / size_of::(); } /// Returns the spare capacity of the buffer as a slice of `MaybeUninit`. @@ -333,13 +575,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 +591,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 +610,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 +638,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 +647,6 @@ impl BufferMut { dst.write(item); dst = dst.add(1); } - self.bytes.set_len(self.bytes.len() + (n * size_of::())); } self.length += n; } @@ -430,73 +666,28 @@ 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); - 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. - /// - /// 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_bytes = self.bytes.split_off(bytes_at); - - // Adjust the lengths, given that length may be < at - let new_length = self.length.saturating_sub(at); - self.length = self.length.min(at); - - BufferMut { - bytes: new_bytes, - length: new_length, - alignment: self.alignment, - _marker: Default::default(), - } - } - - /// 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()). - 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 + // 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.bytes.unsplit(other.bytes); - self.length += other.length; + self.length += slice.len(); } /// Return the [`ByteBufferMut`] for this [`BufferMut`]. pub fn into_byte_buffer(self) -> ByteBufferMut { + let capacity = self + .capacity + .checked_mul(size_of::()) + .vortex_expect("buffer capacity overflow"); ByteBufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length * size_of::(), + capacity, alignment: self.alignment, _marker: Default::default(), } @@ -504,12 +695,8 @@ 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(), - } + let offset = self.ptr.cast::().addr().get() - self.allocation.ptr().addr().get(); + Buffer::from_allocation(self.allocation, offset, self.length, self.alignment) } /// Map each element of the buffer with a closure. @@ -537,14 +724,13 @@ 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 capacity = self.capacity(); + let allocator = self.allocation.allocator().clone(); + let mut aligned = Self::with_capacity_aligned_in(capacity, alignment, allocator); + aligned.extend_from_slice(&self); + aligned } } @@ -568,8 +754,10 @@ impl BufferMut { ); BufferMut { - bytes: self.bytes, + allocation: self.allocation, + ptr: self.ptr.cast(), length: self.length, + capacity: self.capacity, alignment: self.alignment, _marker: std::marker::PhantomData, } @@ -578,14 +766,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 +847,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 +892,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| { @@ -775,147 +973,86 @@ 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 } } -impl Buf for ByteBufferMut { - fn remaining(&self) -> usize { - self.len() - } +#[cfg(test)] +mod test { + use crate::Alignment; + use crate::BufferMut; + use crate::buffer_mut; - fn chunk(&self) -> &[u8] { - self.as_slice() - } + #[test] + fn capacity() { + let mut n = 57; + let mut buf = BufferMut::::with_capacity_aligned(n, Alignment::new(1024)); + assert!(buf.capacity() >= 57); - 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 - ); + while n > 0 { + buf.push(0); + assert!(buf.capacity() >= n); + n -= 1 } - self.bytes.advance(cnt); - self.length -= cnt; - } -} -/// 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 - ); - } - unsafe { self.bytes.advance_mut(cnt) }; - self.length -= cnt; + assert_eq!(buf.alignment(), Alignment::new(1024)); } - #[inline] - fn chunk_mut(&mut self) -> &mut UninitSlice { - self.bytes.chunk_mut() - } + #[test] + fn growth_preserves_alignment_and_values() { + let alignment = Alignment::new(4096); + let mut buffer = BufferMut::::with_capacity_aligned(1, alignment); - 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()); + for value in 0..10_000 { + buffer.push(value); + assert!(alignment.is_offset_aligned(buffer.as_ptr().addr())); } - } - - #[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) + assert_eq!(buffer.as_slice(), (0..10_000).collect::>()); } -} -/// 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"); + #[test] + fn growth_seeds_and_doubles_logical_capacity() { + let alignment = Alignment::new(64); + let mut buffer = BufferMut::::empty_aligned(alignment); - // 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); - } -} + buffer.push(0); + let capacity = buffer.capacity(); + assert_eq!(capacity, Alignment::DEFAULT_ALIGNMENT.as_usize()); -impl Write for ByteBufferMut { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - self.extend_from_slice(buf); - Ok(buf.len()) + buffer.reserve(capacity); + assert_eq!(buffer.capacity(), capacity * 2); } - fn flush(&mut self) -> std::io::Result<()> { - Ok(()) - } -} + #[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(); -#[cfg(test)] -mod test { - use bytes::Buf; - use bytes::BufMut; + buffer.push(u32::MAX); - use crate::Alignment; - use crate::BufferMut; - use crate::ByteBufferMut; - use crate::buffer_mut; + assert_ne!(buffer.as_ptr(), old_ptr); + assert_eq!(&buffer[..capacity], vec![7; capacity]); + assert_eq!(buffer[capacity], u32::MAX); + } #[test] - fn capacity() { - let mut n = 57; - let mut buf = BufferMut::::with_capacity_aligned(n, Alignment::new(1024)); - assert!(buf.capacity() >= 57); + fn raising_logical_alignment_preserves_capacity() { + let buffer = + BufferMut::::with_capacity_preferred_aligned(1, Alignment::of::(), None); + let capacity = buffer.capacity(); - while n > 0 { - buf.push(0); - assert!(buf.capacity() >= n); - n -= 1 - } + let mut buffer = buffer.aligned(Alignment::new(2)); - assert_eq!(buf.alignment(), Alignment::new(1024)); + 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] @@ -994,27 +1131,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; 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()) } } diff --git a/vortex-datafusion/src/convert/scalars.rs b/vortex-datafusion/src/convert/scalars.rs index 7de1fba3da6..c2b1dfc79b2 100644 --- a/vortex-datafusion/src/convert/scalars.rs +++ b/vortex-datafusion/src/convert/scalars.rs @@ -113,13 +113,13 @@ impl TryToDataFusion for Scalar { } // SAFETY: By construction Utf8 scalar values are utf8 DType::Utf8(_) => ScalarValue::Utf8(self.as_utf8().value().cloned().map(|s| unsafe { - String::from_utf8_unchecked(Vec::::from(s.into_inner().into_inner())) + String::from_utf8_unchecked(Vec::::from(s.into_inner().into_bytes())) })), DType::Binary(_) => ScalarValue::Binary( self.as_binary() .value() .cloned() - .map(|b| Vec::::from(b.into_inner())), + .map(|b| Vec::::from(b.into_bytes())), ), dtype @ DType::List(..) => vortex_bail!( "cannot convert Vortex scalar dtype {dtype} to DataFusion ScalarValue: unsupported scalar type" @@ -794,7 +794,7 @@ mod tests { .value() .cloned() .unwrap() - .into_inner() + .into_bytes() .into(); assert_eq!(result_bytes, vec![1u8, 2, 3, 4, 5]); } diff --git a/vortex-ipc/src/messages/encoder.rs b/vortex-ipc/src/messages/encoder.rs index c4372b3cdd9..4b2a0b7007e 100644 --- a/vortex-ipc/src/messages/encoder.rs +++ b/vortex-ipc/src/messages/encoder.rs @@ -81,7 +81,7 @@ impl MessageEncoder { ) .as_union_value(); - buffers.extend(array_buffers.into_iter().map(|b| b.into_inner())); + buffers.extend(array_buffers.into_iter().map(|b| b.into_bytes())); (header, body_len) } @@ -94,7 +94,7 @@ impl MessageEncoder { ) .as_union_value(); let body_len = buffer.len() as u64; - buffers.push(buffer.clone().into_inner()); + buffers.push(buffer.clone().into_bytes()); (header, body_len) } @@ -102,7 +102,7 @@ impl MessageEncoder { let header = fb::DTypeMessage::create(&mut fbb, &fb::DTypeMessageArgs {}).as_union_value(); - let buffer = dtype.write_flatbuffer_bytes()?.into_inner().into_inner(); + let buffer = dtype.write_flatbuffer_bytes()?.into_inner().into_bytes(); let body_len = buffer.len() as u64; buffers.push(buffer); @@ -129,7 +129,7 @@ impl MessageEncoder { .map_err(|_| vortex_err!("Array flatbuffer length must fit into u32"))?; buffers[0] = Bytes::from(fb_buffer_len.to_le_bytes().to_vec()); - buffers[1] = fb_buffer.into_inner().into_inner(); + buffers[1] = fb_buffer.into_inner().into_bytes(); Ok(buffers) }