refactor(array): propagate allocators through builders

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

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()),
         )
     }

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),
         }
     }
 

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::<T>(
+            capacity,
+            decimal,
+            nullability,
+            BufferAllocatorRef::statically_allocated(),
+        )
+    }
+
+    /// Creates a decimal builder with the given capacity and allocator.
+    pub fn with_capacity_in<T: NativeDecimalType>(
+        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::<D>::with_capacity(capacity))
+                DecimalBuffer::from(BufferMut::<D>::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::<T>(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<V: NativeDecimalType>(&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::<i8>::empty())
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use crate::VortexSessionExecute;

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),
         }
     }

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<DType>,
+        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,

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<BitBufferMut>,
     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);
         }

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,16 +81,38 @@ impl<O: OffsetBuilderPType> ListBuilder<O> {
         elements_capacity: usize,
         capacity: usize,
     ) -> Self {
-        let elements_builder = ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity);
-        let mut offsets_builder = PrimitiveBuilder::<O>::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<DType>,
+        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::<O>::with_capacity_in(NonNullable, capacity + 1, allocator.clone());
 
         // The first offset is always 0 and represents an empty list.
         offsets_builder.append_zero();
 
         Self {
             elements_builder,
             offsets_builder,
-            nulls: ValidityBuilder::new(capacity),
+            nulls: ValidityBuilder::new_in(capacity, allocator),
             dtype: DType::List(value_dtype, nullability),
         }
     }

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<O: OffsetBuilderPType, S: OffsetBuilderPType> ListViewBuilder<O, S> {
         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<DType>,
+        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::<O>::with_capacity(Nullability::NonNullable, capacity);
-        let sizes_builder =
-            PrimitiveBuilder::<S>::with_capacity(Nullability::NonNullable, capacity);
+        let offsets_builder = PrimitiveBuilder::<O>::with_capacity_in(
+            Nullability::NonNullable,
+            capacity,
+            allocator.clone(),
+        );
+        let sizes_builder = PrimitiveBuilder::<S>::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),