diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 3587f6fc3ba..9f09cbf78b1 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -5,6 +5,7 @@ use std::any::Any; use std::mem; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -33,9 +34,22 @@ impl BoolBuilder { } pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self { + Self::with_capacity_in( + nullability, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a builder with the given capacity and allocator. + pub fn with_capacity_in( + nullability: Nullability, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { Self { - inner: BitBufferMut::with_capacity(capacity), - nulls: LazyBitBufferBuilder::new(capacity), + inner: BitBufferMut::with_capacity_in(capacity, allocator.clone()), + nulls: LazyBitBufferBuilder::new_in(capacity, allocator), dtype: DType::Bool(nullability), } } @@ -61,8 +75,10 @@ impl BoolBuilder { "Null count and value count should match when calling BoolBuilder::finish." ); + let allocator = self.inner.allocator().clone(); + let inner = mem::replace(&mut self.inner, BitBufferMut::empty_in(allocator)).freeze(); BoolArray::new( - mem::take(&mut self.inner).freeze(), + inner, self.nulls.finish_with_nullability(self.dtype.nullability()), ) } diff --git a/vortex-array/src/builders/child.rs b/vortex-array/src/builders/child.rs index d1ff5f2954e..da047404405 100644 --- a/vortex-array/src/builders/child.rs +++ b/vortex-array/src/builders/child.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -9,7 +10,7 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ChunkedArray; use crate::builders::ArrayBuilder; -use crate::builders::builder_with_capacity; +use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::scalar::Scalar; @@ -42,12 +43,18 @@ pub struct ChildBuilder { impl ChildBuilder { /// Creates a new `ChildBuilder` whose scalar builder is pre-allocated for `capacity` values. + #[cfg(test)] pub fn with_capacity(dtype: &DType, capacity: usize) -> Self { + Self::with_capacity_in(BufferAllocatorRef::statically_allocated(), dtype, capacity) + } + + /// Creates a child builder with the provided allocator. + pub fn with_capacity_in(allocator: BufferAllocatorRef, dtype: &DType, capacity: usize) -> Self { Self { dtype: dtype.clone(), chunks: Vec::new(), chunks_len: 0, - pending: builder_with_capacity(dtype, capacity), + pending: builder_with_capacity_in(allocator, dtype, capacity), } } diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index 027b7e62c75..0c33247b980 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -3,6 +3,7 @@ use std::any::Any; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -103,13 +104,31 @@ impl DecimalBuilder { capacity: usize, decimal: DecimalDType, nullability: Nullability, + ) -> Self { + Self::with_capacity_in::( + capacity, + decimal, + nullability, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a decimal builder with the given capacity and allocator. + pub fn with_capacity_in( + capacity: usize, + decimal: DecimalDType, + nullability: Nullability, + allocator: BufferAllocatorRef, ) -> Self { Self { dtype: DType::Decimal(decimal, nullability), values: match_each_decimal_value_type!(T::DECIMAL_TYPE, |D| { - DecimalBuffer::from(BufferMut::::with_capacity(capacity)) + DecimalBuffer::from(BufferMut::::with_capacity_in( + capacity, + allocator.clone(), + )) }), - nulls: LazyBitBufferBuilder::new(capacity), + nulls: LazyBitBufferBuilder::new_in(capacity, allocator), } } @@ -153,7 +172,7 @@ impl DecimalBuilder { let decimal_dtype = *self.decimal_dtype(); - delegate_fn!(std::mem::take(&mut self.values), |T, values| { + delegate_fn!(self.values.take(), |T, values| { DecimalArray::new::(values.freeze(), decimal_dtype, validity) }) } @@ -228,6 +247,13 @@ impl ArrayBuilder for DecimalBuilder { } impl DecimalBuffer { + fn take(&mut self) -> Self { + delegate_fn!(self, |T, buffer| { + let allocator = buffer.allocator(); + DecimalBuffer::from(std::mem::replace(buffer, allocator.with_capacity(0))) + }) + } + fn push(&mut self, value: V) { delegate_fn!(self, |T, buffer| { buffer.push( @@ -291,12 +317,6 @@ impl_from_buffer!(i64, I64); impl_from_buffer!(i128, I128); impl_from_buffer!(i256, I256); -impl Default for DecimalBuffer { - fn default() -> Self { - Self::I8(BufferMut::::empty()) - } -} - #[cfg(test)] mod tests { use crate::VortexSessionExecute; diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index c1a91f202d0..81af4fcc1b7 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -3,6 +3,7 @@ use std::any::Any; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -34,8 +35,21 @@ impl ExtensionBuilder { /// Creates a new `ExtensionBuilder` with the given `capacity`. pub fn with_capacity(ext_dtype: ExtDTypeRef, capacity: usize) -> Self { + Self::with_capacity_in( + ext_dtype, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates an extension builder with the provided allocator. + pub fn with_capacity_in( + ext_dtype: ExtDTypeRef, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { Self { - storage: ChildBuilder::with_capacity(ext_dtype.storage_dtype(), capacity), + storage: ChildBuilder::with_capacity_in(allocator, ext_dtype.storage_dtype(), capacity), dtype: DType::Extension(ext_dtype), } } diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 98acd2c468e..5b2efd0f5dc 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -59,12 +60,30 @@ impl FixedSizeListBuilder { list_size: u32, nullability: Nullability, capacity: usize, + ) -> Self { + Self::with_capacity_in( + element_dtype, + list_size, + nullability, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a fixed-size-list builder with the provided allocator. + pub fn with_capacity_in( + element_dtype: Arc, + list_size: u32, + nullability: Nullability, + capacity: usize, + allocator: BufferAllocatorRef, ) -> Self { let elements_capacity = capacity * list_size as usize; - let elements_builder = ChildBuilder::with_capacity(&element_dtype, elements_capacity); + let elements_builder = + ChildBuilder::with_capacity_in(allocator.clone(), &element_dtype, elements_capacity); let fsl_dtype = DType::FixedSizeList(element_dtype, list_size, nullability); - let nulls = ValidityBuilder::new(capacity); + let nulls = ValidityBuilder::new_in(capacity, allocator); Self { dtype: fsl_dtype, diff --git a/vortex-array/src/builders/lazy_null_builder.rs b/vortex-array/src/builders/lazy_null_builder.rs index 8a9f62d1a03..b3b9be18015 100644 --- a/vortex-array/src/builders/lazy_null_builder.rs +++ b/vortex-array/src/builders/lazy_null_builder.rs @@ -3,6 +3,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -18,16 +19,23 @@ pub struct LazyBitBufferBuilder { inner: Option, len: usize, capacity: usize, + allocator: BufferAllocatorRef, } impl LazyBitBufferBuilder { /// Creates a new empty builder. /// `capacity` is the number of bits in the null buffer. pub fn new(capacity: usize) -> Self { + Self::new_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Creates a new empty builder with the provided allocator. + pub fn new_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { Self { inner: None, len: 0, capacity, + allocator, } } @@ -148,7 +156,8 @@ impl LazyBitBufferBuilder { #[inline(never)] fn materialize(&mut self) { if self.inner.is_none() { - let mut bit_mut = BitBufferMut::with_capacity(self.len.max(self.capacity)); + let mut bit_mut = + BitBufferMut::with_capacity_in(self.len.max(self.capacity), self.allocator.clone()); bit_mut.append_n(true, self.len); self.inner = Some(bit_mut); } diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index ca74b7b9328..5b10981e25b 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -5,6 +5,7 @@ use std::any::Any; use std::sync::Arc; use num_traits::AsPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -80,8 +81,30 @@ impl ListBuilder { elements_capacity: usize, capacity: usize, ) -> Self { - let elements_builder = ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity); - let mut offsets_builder = PrimitiveBuilder::::with_capacity(NonNullable, capacity + 1); + Self::with_capacity_in( + value_dtype, + nullability, + elements_capacity, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a list builder with the provided allocator. + pub fn with_capacity_in( + value_dtype: Arc, + nullability: Nullability, + elements_capacity: usize, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { + let elements_builder = ChildBuilder::with_capacity_in( + allocator.clone(), + value_dtype.as_ref(), + elements_capacity, + ); + let mut offsets_builder = + PrimitiveBuilder::::with_capacity_in(NonNullable, capacity + 1, allocator.clone()); // The first offset is always 0 and represents an empty list. offsets_builder.append_zero(); @@ -89,7 +112,7 @@ impl ListBuilder { Self { elements_builder, offsets_builder, - nulls: ValidityBuilder::new(capacity), + nulls: ValidityBuilder::new_in(capacity, allocator), dtype: DType::List(value_dtype, nullability), } } diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 2efc2d45d74..57a7c29af78 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use num_traits::ToPrimitive; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -104,14 +105,38 @@ impl ListViewBuilder { elements_capacity: usize, capacity: usize, ) -> Self { - let elements_builder = ChildBuilder::with_capacity(&element_dtype, elements_capacity); + Self::with_capacity_in( + element_dtype, + nullability, + elements_capacity, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a list-view builder with the provided allocator. + pub fn with_capacity_in( + element_dtype: Arc, + nullability: Nullability, + elements_capacity: usize, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { + let elements_builder = + ChildBuilder::with_capacity_in(allocator.clone(), &element_dtype, elements_capacity); - let offsets_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); - let sizes_builder = - PrimitiveBuilder::::with_capacity(Nullability::NonNullable, capacity); + let offsets_builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + capacity, + allocator.clone(), + ); + let sizes_builder = PrimitiveBuilder::::with_capacity_in( + Nullability::NonNullable, + capacity, + allocator.clone(), + ); - let nulls = ValidityBuilder::new(capacity); + let nulls = ValidityBuilder::new_in(capacity, allocator); Self { dtype: DType::List(element_dtype, nullability), diff --git a/vortex-array/src/builders/map.rs b/vortex-array/src/builders/map.rs index 6714cf8f881..09b36a3fe2d 100644 --- a/vortex-array/src/builders/map.rs +++ b/vortex-array/src/builders/map.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -45,11 +46,27 @@ impl MapBuilder { /// Creates a map builder with space for `capacity` map rows. pub fn with_capacity(map_dtype: MapDType, nullability: Nullability, capacity: usize) -> Self { - let entries_builder = ListViewBuilder::with_capacity( + Self::with_capacity_in( + map_dtype, + nullability, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a map builder with the provided allocator. + pub fn with_capacity_in( + map_dtype: MapDType, + nullability: Nullability, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { + let entries_builder = ListViewBuilder::with_capacity_in( Arc::new(map_dtype.entries_dtype()), nullability, capacity.saturating_mul(2), capacity, + allocator, ); let dtype = DType::Map(map_dtype.clone(), nullability); Self { diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index ea8ee4d74a0..254e423d8e3 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -387,72 +387,84 @@ macro_rules! __match_each_map_builder_size { /// assert_eq!(strings.execute_scalar(3, &mut ctx).unwrap(), "d".into()); /// ``` pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box { + builder_with_capacity_in(BufferAllocatorRef::statically_allocated(), dtype, capacity) +} + +/// Construct a canonical builder using the provided buffer allocator. +pub fn builder_with_capacity_in( + allocator: BufferAllocatorRef, + dtype: &DType, + capacity: usize, +) -> Box { match dtype { DType::Null => Box::new(NullBuilder::new()), - DType::Bool(n) => Box::new(BoolBuilder::with_capacity(*n, capacity)), + DType::Bool(n) => Box::new(BoolBuilder::with_capacity_in(*n, capacity, allocator)), DType::Primitive(ptype, n) => { match_each_native_ptype!(ptype, |P| { - Box::new(PrimitiveBuilder::

::with_capacity(*n, capacity)) + Box::new(PrimitiveBuilder::

::with_capacity_in( + *n, capacity, allocator, + )) }) } DType::Decimal(decimal_type, n) => { match_each_decimal_value_type!( DecimalType::smallest_decimal_value_type(decimal_type), |D| { - Box::new(DecimalBuilder::with_capacity::( + Box::new(DecimalBuilder::with_capacity_in::( capacity, *decimal_type, *n, + allocator, )) } ) } - DType::Utf8(n) => Box::new(VarBinViewBuilder::with_capacity(DType::Utf8(*n), capacity)), - DType::Binary(n) => Box::new(VarBinViewBuilder::with_capacity( + DType::Utf8(n) => Box::new(VarBinViewBuilder::with_capacity_in( + DType::Utf8(*n), + capacity, + allocator, + )), + DType::Binary(n) => Box::new(VarBinViewBuilder::with_capacity_in( DType::Binary(*n), capacity, + allocator, )), - DType::List(dtype, n) => Box::new(ListViewBuilder::::with_capacity( + DType::List(dtype, n) => Box::new(ListViewBuilder::::with_capacity_in( Arc::clone(dtype), *n, 2 * capacity, // Arbitrarily choose 2 times the `offsets` capacity here. capacity, + allocator, )), - DType::Map(map_dtype, nullability) => Box::new(MapBuilder::::with_capacity( + DType::Map(map_dtype, nullability) => Box::new(MapBuilder::::with_capacity_in( map_dtype.clone(), *nullability, capacity, + allocator, )), DType::FixedSizeList(elem_dtype, list_size, null) => { - Box::new(FixedSizeListBuilder::with_capacity( + Box::new(FixedSizeListBuilder::with_capacity_in( Arc::clone(elem_dtype), *list_size, *null, capacity, + allocator, )) } - DType::Struct(struct_dtype, n) => Box::new(StructBuilder::with_capacity( + DType::Struct(struct_dtype, n) => Box::new(StructBuilder::with_capacity_in( struct_dtype.clone(), *n, capacity, + allocator, )), DType::Union(..) => todo!("TODO(connor)[Union]: unimplemented"), DType::Variant(_) => { unimplemented!() } - DType::Extension(ext_dtype) => { - Box::new(ExtensionBuilder::with_capacity(ext_dtype.clone(), capacity)) - } + DType::Extension(ext_dtype) => Box::new(ExtensionBuilder::with_capacity_in( + ext_dtype.clone(), + capacity, + allocator, + )), } } - -/// Construct a new canonical builder for the given [`DType`] using a host -/// [`vortex_buffer::BufferAllocator`]. -pub fn builder_with_capacity_in( - allocator: BufferAllocatorRef, - dtype: &DType, - capacity: usize, -) -> Box { - let _allocator = allocator; - builder_with_capacity(dtype, capacity) -} diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index aca2db36286..a8481dca79e 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::mem::MaybeUninit; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -38,9 +39,22 @@ impl PrimitiveBuilder { /// Creates a new `PrimitiveBuilder` with the given `capacity`. pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self { + Self::with_capacity_in( + nullability, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a builder with the given capacity and allocator. + pub fn with_capacity_in( + nullability: Nullability, + capacity: usize, + allocator: BufferAllocatorRef, + ) -> Self { Self { - values: BufferMut::with_capacity(capacity), - nulls: LazyBitBufferBuilder::new(capacity), + values: BufferMut::with_capacity_in(capacity, allocator.clone()), + nulls: LazyBitBufferBuilder::new_in(capacity, allocator), dtype: DType::Primitive(T::PTYPE, nullability), } } @@ -120,7 +134,9 @@ impl PrimitiveBuilder { .nulls .finish_with_nullability(self.dtype().nullability()); - PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity) + let allocator = self.values.allocator().clone(); + let values = std::mem::replace(&mut self.values, allocator.with_capacity(0)).freeze(); + PrimitiveArray::new(values, validity) } /// Extends the primitive array with an iterator. diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 66a2c66358d..d4a558cc9b0 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -4,6 +4,7 @@ use std::any::Any; use itertools::Itertools; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -44,15 +45,30 @@ impl StructBuilder { struct_dtype: StructFields, nullability: Nullability, capacity: usize, + ) -> Self { + Self::with_capacity_in( + struct_dtype, + nullability, + capacity, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a struct builder with the provided allocator. + pub fn with_capacity_in( + struct_dtype: StructFields, + nullability: Nullability, + capacity: usize, + allocator: BufferAllocatorRef, ) -> Self { let builders = struct_dtype .fields() - .map(|dt| ChildBuilder::with_capacity(&dt, capacity)) + .map(|dt| ChildBuilder::with_capacity_in(allocator.clone(), &dt, capacity)) .collect(); Self { builders, - nulls: ValidityBuilder::new(capacity), + nulls: ValidityBuilder::new_in(capacity, allocator), dtype: DType::Struct(struct_dtype, nullability), } } diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 914cd04e1b1..f4922c3bcdf 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -1,10 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +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 rstest::rstest; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -42,6 +50,7 @@ use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::ListBuilder; use crate::builders::builder_with_capacity; +use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; @@ -53,6 +62,56 @@ use crate::extension::datetime::Timestamp; use crate::scalar::Scalar; use crate::validity::Validity; +#[derive(Debug)] +struct CountingAllocator { + allocations: Arc, +} + +// SAFETY: this forwards memory operations to Global and only counts allocations. +unsafe impl Allocator for CountingAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + self.allocations.fetch_add(1, Ordering::Relaxed); + Global.allocate(layout) + } + + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // SAFETY: ptr and layout came from Global. + unsafe { Global.deallocate(ptr, layout) } + } +} + +#[test] +fn builder_reuses_its_allocator_after_finish() { + let allocations = Arc::new(AtomicUsize::new(0)); + let allocator = BufferAllocatorRef::new(CountingAllocator { + allocations: Arc::clone(&allocations), + }); + let dtype = DType::Struct( + StructFields::from_iter([ + ( + "number", + DType::Primitive(PType::I32, Nullability::Nullable), + ), + ("text", DType::Utf8(Nullability::Nullable)), + ]), + Nullability::Nullable, + ); + let mut builder = builder_with_capacity_in(allocator, &dtype, 1); + + builder.append_null(); + drop(builder.finish()); + let first = allocations.load(Ordering::Relaxed); + + builder.append_null(); + drop(builder.finish()); + let second = allocations.load(Ordering::Relaxed); + + assert!( + second - first >= 5, + "all nested buffers must reuse the allocator" + ); +} + /// Test that `append_zeros` produces the same result as manually appending `Scalar::default_value`. /// /// This test verifies that the implementation of `append_zeros` correctly matches the behavior diff --git a/vortex-array/src/builders/validity.rs b/vortex-array/src/builders/validity.rs index d5417407b2f..31832486b74 100644 --- a/vortex-array/src/builders/validity.rs +++ b/vortex-array/src/builders/validity.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexExpect; use vortex_error::vortex_panic; @@ -37,11 +38,17 @@ pub(crate) struct ValidityBuilder { impl ValidityBuilder { /// Creates a new `ValidityBuilder` whose null buffer is pre-allocated for `capacity` bits. + #[cfg(test)] pub fn new(capacity: usize) -> Self { + Self::new_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Creates a validity builder with the provided allocator. + pub fn new_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { Self { runs: Vec::new(), runs_len: 0, - pending: LazyBitBufferBuilder::new(capacity), + pending: LazyBitBufferBuilder::new_in(capacity, allocator), } } diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 3a32c2e9e6b..76b9999b0d6 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -9,6 +9,7 @@ use itertools::Itertools; use num_traits::AsPrimitive; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -46,11 +47,24 @@ pub struct VarBinViewBuilder { in_progress: Option, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, + allocator: BufferAllocatorRef, } impl VarBinViewBuilder { pub fn with_capacity(dtype: DType, capacity: usize) -> Self { - Self::new(dtype, capacity, Default::default(), Default::default(), 0.0) + Self::with_capacity_in(dtype, capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Creates a builder with the given capacity and allocator. + pub fn with_capacity_in(dtype: DType, capacity: usize, allocator: BufferAllocatorRef) -> Self { + Self::new_in( + dtype, + capacity, + Default::default(), + Default::default(), + 0.0, + allocator, + ) } pub fn with_buffer_deduplication(dtype: DType, capacity: usize) -> Self { @@ -79,23 +93,44 @@ impl VarBinViewBuilder { completed: CompletedBuffers, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, + ) -> Self { + Self::new_in( + dtype, + capacity, + completed, + growth_strategy, + compaction_threshold, + BufferAllocatorRef::statically_allocated(), + ) + } + + /// Creates a configurable builder with the provided allocator. + pub fn new_in( + dtype: DType, + capacity: usize, + completed: CompletedBuffers, + growth_strategy: BufferGrowthStrategy, + compaction_threshold: f64, + allocator: BufferAllocatorRef, ) -> Self { assert!( matches!(dtype, DType::Utf8(_) | DType::Binary(_)), "VarBinViewBuilder DType must be Utf8 or Binary." ); Self { - views_builder: BufferMut::with_capacity_preferred_aligned( + views_builder: BufferMut::with_capacity_preferred_aligned_in( capacity, Alignment::of::(), None, + allocator.clone(), ), - nulls: LazyBitBufferBuilder::new(capacity), + nulls: LazyBitBufferBuilder::new_in(capacity, allocator.clone()), completed, in_progress: None, dtype, growth_strategy, compaction_threshold, + allocator, } } @@ -153,10 +188,11 @@ impl VarBinViewBuilder { fn init_in_progress(&mut self, min_len: usize) { let next_buffer_size = self.growth_strategy.next_size() as usize; let to_reserve = next_buffer_size.max(min_len); - self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned( + self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned_in( to_reserve, Alignment::of::(), None, + self.allocator.clone(), )); } @@ -616,7 +652,7 @@ impl VarBinViewBuilder { referenced: usize, ) { let buf_index = self.completed.len(); - let mut compact = ByteBufferMut::with_capacity(referenced); + let mut compact = ByteBufferMut::with_capacity_in(referenced, self.allocator.clone()); self.views_builder.reserve(count); let data = bytes.as_slice(); @@ -667,14 +703,15 @@ impl VarBinViewBuilder { let validity = self.nulls.finish_with_nullability(self.dtype.nullability()); + let views = std::mem::replace( + &mut self.views_builder, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + // SAFETY: the builder methods check safety at each step. unsafe { - VarBinViewArray::new_unchecked( - std::mem::take(&mut self.views_builder).freeze(), - buffers.finish(), - self.dtype.clone(), - validity, - ) + VarBinViewArray::new_unchecked(views, buffers.finish(), self.dtype.clone(), validity) } } diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index 734f65a08a6..c0c6338378d 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -6,6 +6,7 @@ use std::ops::Not; use bitvec::view::BitView; use crate::BitBuffer; +use crate::BufferAllocatorRef; use crate::BufferMut; use crate::ByteBufferMut; use crate::bit::collect_bool_words; @@ -118,8 +119,13 @@ impl BitBufferMut { /// Creates a `BitBufferMut` from a [`BitBuffer`] by copying all of the data over. pub fn copy_from(bit_buffer: &BitBuffer) -> Self { + Self::copy_from_in(bit_buffer, bit_buffer.inner().allocator().clone()) + } + + /// Copies a bit buffer with the provided allocator. + pub fn copy_from_in(bit_buffer: &BitBuffer, allocator: BufferAllocatorRef) -> Self { Self { - buffer: ByteBufferMut::copy_from(bit_buffer.inner()), + buffer: ByteBufferMut::copy_from_in(bit_buffer.inner(), allocator), offset: bit_buffer.offset(), len: bit_buffer.len(), } @@ -128,8 +134,14 @@ impl BitBufferMut { /// Create a new empty mutable bit buffer with requested capacity (in bits). #[inline] pub fn with_capacity(capacity: usize) -> Self { + Self::with_capacity_in(capacity, BufferAllocatorRef::statically_allocated()) + } + + /// Create a mutable bit buffer with the provided allocator. + #[inline] + pub fn with_capacity_in(capacity: usize, allocator: BufferAllocatorRef) -> Self { Self { - buffer: BufferMut::with_capacity(capacity.div_ceil(8)), + buffer: BufferMut::with_capacity_in(capacity.div_ceil(8), allocator), offset: 0, len: 0, } @@ -161,6 +173,17 @@ impl BitBufferMut { Self::with_capacity(0) } + /// Create an empty mutable bit buffer with the provided allocator. + #[inline(always)] + pub fn empty_in(allocator: BufferAllocatorRef) -> Self { + Self::with_capacity_in(0, allocator) + } + + /// Returns the allocator that owns this buffer. + pub fn allocator(&self) -> &BufferAllocatorRef { + self.buffer.allocator() + } + /// Create a new mutable buffer with requested `len` and all bits set to `value`. #[inline] pub fn full(value: bool, len: usize) -> Self {