Buffer Allocator Abstraction (#7337)

Initial PR to move towards a session-defined buffer allocation.

1. Sets up a MemorySession (we could move this and all buffer handle
types to vortex-io and treat the i/o crate as engine integration?)
2. Defines an allocator trait, returning a WritableHostBuffer
3. Freeze the WriteableHostBuffer into a host ByteBuffer

For now, we keep the existing behavior for device buffers (i.e.
vortex-cuda constructs its own device BufferHandle inside its segment
source).

This change configures host VortexReadAt sources (ObjectStoreReadAt and
FileReadAt) to use the allocator from the session. It also demonstrates
using the allocator from within an array (chunked canonical), although
it doesn't do the plumbing work all the way into the array builders.

A follow up PR will implement a pooling allocator.

Further PRs wlll port over the execution logic to use the allocator from
the ExecutionCtx, eventually using clippy to forbid direct ByteBufferMut
access from within our (non allocator) crates.

---------

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

Author: gatesn

vortex-array/public-api.lock

@@ -8388,6 +8388,8 @@ pub unsafe fn vortex_array::builders::PrimitiveBuilder<T>::set_validity_unchecke
 
 pub fn vortex_array::builders::builder_with_capacity(dtype: &vortex_array::dtype::DType, capacity: usize) -> alloc::boxed::Box<dyn vortex_array::builders::ArrayBuilder>
 
+pub fn vortex_array::builders::builder_with_capacity_in(allocator: vortex_array::memory::HostAllocatorRef, dtype: &vortex_array::dtype::DType, capacity: usize) -> alloc::boxed::Box<dyn vortex_array::builders::ArrayBuilder>
+
 pub mod vortex_array::builtins
 
 pub trait vortex_array::builtins::ArrayBuiltins: core::marker::Sized
@@ -13262,6 +13264,110 @@ pub fn V::matches(array: &vortex_array::ArrayRef) -> bool
 
 pub fn V::try_match<'a>(array: &'a vortex_array::ArrayRef) -> core::option::Option<vortex_array::ArrayView<'a, V>>
 
+pub mod vortex_array::memory
+
+pub struct vortex_array::memory::DefaultHostAllocator
+
+impl core::default::Default for vortex_array::memory::DefaultHostAllocator
+
+pub fn vortex_array::memory::DefaultHostAllocator::default() -> vortex_array::memory::DefaultHostAllocator
+
+impl core::fmt::Debug for vortex_array::memory::DefaultHostAllocator
+
+pub fn vortex_array::memory::DefaultHostAllocator::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::memory::HostAllocator for vortex_array::memory::DefaultHostAllocator
+
+pub fn vortex_array::memory::DefaultHostAllocator::allocate(&self, len: usize, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+pub struct vortex_array::memory::MemorySession
+
+impl vortex_array::memory::MemorySession
+
+pub fn vortex_array::memory::MemorySession::allocator(&self) -> vortex_array::memory::HostAllocatorRef
+
+pub fn vortex_array::memory::MemorySession::new(allocator: vortex_array::memory::HostAllocatorRef) -> Self
+
+pub fn vortex_array::memory::MemorySession::set_allocator(&mut self, allocator: vortex_array::memory::HostAllocatorRef)
+
+impl core::default::Default for vortex_array::memory::MemorySession
+
+pub fn vortex_array::memory::MemorySession::default() -> Self
+
+impl core::fmt::Debug for vortex_array::memory::MemorySession
+
+pub fn vortex_array::memory::MemorySession::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_array::memory::WritableHostBuffer
+
+impl vortex_array::memory::WritableHostBuffer
+
+pub fn vortex_array::memory::WritableHostBuffer::alignment(&self) -> vortex_buffer::alignment::Alignment
+
+pub fn vortex_array::memory::WritableHostBuffer::as_mut_slice(&mut self) -> &mut [u8]
+
+pub fn vortex_array::memory::WritableHostBuffer::as_mut_slice_typed<T>(&mut self) -> vortex_error::VortexResult<&mut [T]>
+
+pub fn vortex_array::memory::WritableHostBuffer::freeze(self) -> vortex_buffer::ByteBuffer
+
+pub fn vortex_array::memory::WritableHostBuffer::freeze_typed<T>(self) -> vortex_error::VortexResult<vortex_buffer::buffer::Buffer<T>>
+
+pub fn vortex_array::memory::WritableHostBuffer::is_empty(&self) -> bool
+
+pub fn vortex_array::memory::WritableHostBuffer::len(&self) -> usize
+
+pub fn vortex_array::memory::WritableHostBuffer::new(inner: alloc::boxed::Box<dyn vortex_array::memory::HostBufferMut>) -> Self
+
+impl core::fmt::Debug for vortex_array::memory::WritableHostBuffer
+
+pub fn vortex_array::memory::WritableHostBuffer::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub trait vortex_array::memory::HostAllocator: core::fmt::Debug + core::marker::Send + core::marker::Sync + 'static
+
+pub fn vortex_array::memory::HostAllocator::allocate(&self, len: usize, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+impl vortex_array::memory::HostAllocator for vortex_array::memory::DefaultHostAllocator
+
+pub fn vortex_array::memory::DefaultHostAllocator::allocate(&self, len: usize, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+pub trait vortex_array::memory::HostAllocatorExt: vortex_array::memory::HostAllocator
+
+pub fn vortex_array::memory::HostAllocatorExt::allocate_typed<T>(&self, len: usize) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+impl<A: vortex_array::memory::HostAllocator + ?core::marker::Sized> vortex_array::memory::HostAllocatorExt for A
+
+pub fn A::allocate_typed<T>(&self, len: usize) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+pub trait vortex_array::memory::HostBufferMut: core::marker::Send + 'static
+
+pub fn vortex_array::memory::HostBufferMut::alignment(&self) -> vortex_buffer::alignment::Alignment
+
+pub fn vortex_array::memory::HostBufferMut::as_mut_slice(&mut self) -> &mut [u8]
+
+pub fn vortex_array::memory::HostBufferMut::freeze(self: alloc::boxed::Box<Self>) -> vortex_buffer::ByteBuffer
+
+pub fn vortex_array::memory::HostBufferMut::is_empty(&self) -> bool
+
+pub fn vortex_array::memory::HostBufferMut::len(&self) -> usize
+
+pub trait vortex_array::memory::MemorySessionExt: vortex_session::SessionExt
+
+pub fn vortex_array::memory::MemorySessionExt::allocator(&self) -> vortex_array::memory::HostAllocatorRef
+
+pub fn vortex_array::memory::MemorySessionExt::memory(&self) -> vortex_session::Ref<'_, vortex_array::memory::MemorySession>
+
+pub fn vortex_array::memory::MemorySessionExt::memory_mut(&self) -> vortex_session::RefMut<'_, vortex_array::memory::MemorySession>
+
+impl<S: vortex_session::SessionExt> vortex_array::memory::MemorySessionExt for S
+
+pub fn S::allocator(&self) -> vortex_array::memory::HostAllocatorRef
+
+pub fn S::memory(&self) -> vortex_session::Ref<'_, vortex_array::memory::MemorySession>
+
+pub fn S::memory_mut(&self) -> vortex_session::RefMut<'_, vortex_array::memory::MemorySession>
+
+pub type vortex_array::memory::HostAllocatorRef = alloc::sync::Arc<dyn vortex_array::memory::HostAllocator>
+
 pub mod vortex_array::normalize
 
 pub enum vortex_array::normalize::Operation<'a>
@@ -22362,6 +22468,8 @@ pub struct vortex_array::ExecutionCtx
 
 impl vortex_array::ExecutionCtx
 
+pub fn vortex_array::ExecutionCtx::allocator(&self) -> vortex_array::memory::HostAllocatorRef
+
 pub fn vortex_array::ExecutionCtx::log(&mut self, msg: core::fmt::Arguments<'_>)
 
 pub fn vortex_array::ExecutionCtx::new(session: vortex_session::VortexSession) -> Self

vortex-array/src/arrays/chunked/vtable/canonical.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_buffer::BufferMut;
+use vortex_buffer::Buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
@@ -19,12 +19,13 @@ use crate::arrays::chunked::ChunkedArrayExt;
 use crate::arrays::listview::ListViewArrayExt;
 use crate::arrays::listview::ListViewRebuildMode;
 use crate::arrays::struct_::StructArrayExt;
-use crate::builders::builder_with_capacity;
+use crate::builders::builder_with_capacity_in;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::dtype::StructFields;
+use crate::memory::HostAllocatorExt;
 use crate::validity::Validity;
 
 pub(super) fn _canonicalize(
@@ -56,7 +57,7 @@ pub(super) fn _canonicalize(
             ctx,
         )?),
         _ => {
-            let mut builder = builder_with_capacity(array.dtype(), array.len());
+            let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len());
             array.array().append_to_builder(builder.as_mut(), ctx)?;
             builder.finish_into_canonical()
         }
@@ -131,8 +132,12 @@ fn swizzle_list_chunks(
     // this much more complicated.
     // We (somewhat arbitrarily) choose `u64` for our offsets and sizes here. These can always be
     // narrowed later by the compressor.
-    let mut offsets = BufferMut::<u64>::with_capacity(len);
-    let mut sizes = BufferMut::<u64>::with_capacity(len);
+    let allocator = ctx.allocator();
+    let mut offsets = allocator.allocate_typed::<u64>(len)?;
+    let mut sizes = allocator.allocate_typed::<u64>(len)?;
+    let offsets_out: &mut [u64] = offsets.as_mut_slice_typed::<u64>()?;
+    let sizes_slice_out: &mut [u64] = sizes.as_mut_slice_typed::<u64>()?;
+    let mut next_list = 0usize;
 
     for chunk in chunks {
         let chunk_array = chunk.clone().execute::<ListViewArray>(ctx)?;
@@ -162,19 +167,31 @@ fn swizzle_list_chunks(
         let sizes_slice = sizes_arr.as_slice::<u64>();
 
         // Append offsets and sizes, adjusting offsets to point into the combined array.
-        offsets.extend(offsets_slice.iter().map(|o| o + num_elements));
-        sizes.extend(sizes_slice);
+        for (&offset, &size) in offsets_slice.iter().zip(sizes_slice.iter()) {
+            offsets_out[next_list] = offset + num_elements;
+            sizes_slice_out[next_list] = size;
+            next_list += 1;
+        }
 
         num_elements += chunk_array.elements().len() as u64;
     }
+    debug_assert_eq!(next_list, len);
 
     // SAFETY: elements are sliced from valid `ListViewArray`s (from `to_listview()`).
     let chunked_elements =
         unsafe { ChunkedArray::new_unchecked(list_elements_chunks, elem_dtype.clone()) }
             .into_array();
 
-    let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable).into_array();
-    let sizes = PrimitiveArray::new(sizes.freeze(), Validity::NonNullable).into_array();
+    let offsets = PrimitiveArray::new(
+        Buffer::<u64>::from_byte_buffer(offsets.freeze()),
+        Validity::NonNullable,
+    )
+    .into_array();
+    let sizes = PrimitiveArray::new(
+        Buffer::<u64>::from_byte_buffer(sizes.freeze()),
+        Validity::NonNullable,
+    )
+    .into_array();
 
     // SAFETY:
     // - `offsets` and `sizes` are non-nullable u64 arrays of the same length
@@ -192,9 +209,15 @@ fn swizzle_list_chunks(
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::atomic::Ordering;
 
     use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
 
+    use crate::Canonical;
+    use crate::ExecutionCtx;
     use crate::IntoArray;
     use crate::ToCanonical;
     use crate::accessor::ArrayAccessor;
@@ -207,8 +230,28 @@ mod tests {
     use crate::dtype::DType::Primitive;
     use crate::dtype::Nullability::NonNullable;
     use crate::dtype::PType::I32;
+    use crate::memory::DefaultHostAllocator;
+    use crate::memory::HostAllocator;
+    use crate::memory::MemorySessionExt;
+    use crate::memory::WritableHostBuffer;
     use crate::validity::Validity;
 
+    #[derive(Debug)]
+    struct CountingAllocator {
+        allocations: Arc<AtomicUsize>,
+    }
+
+    impl HostAllocator for CountingAllocator {
+        fn allocate(
+            &self,
+            len: usize,
+            alignment: vortex_buffer::Alignment,
+        ) -> VortexResult<WritableHostBuffer> {
+            self.allocations.fetch_add(1, Ordering::Relaxed);
+            DefaultHostAllocator.allocate(len, alignment)
+        }
+    }
+
     #[test]
     pub fn pack_nested_structs() {
         let struct_array = StructArray::try_new(
@@ -265,4 +308,42 @@ mod tests {
         assert_eq!(l1.scalar_at(0).unwrap(), canon_values.scalar_at(0).unwrap());
         assert_eq!(l2.scalar_at(0).unwrap(), canon_values.scalar_at(1).unwrap());
     }
+
+    #[test]
+    fn list_canonicalize_uses_memory_session_allocator() {
+        let allocations = Arc::new(AtomicUsize::new(0));
+        let session = VortexSession::empty();
+        session
+            .memory_mut()
+            .set_allocator(Arc::new(CountingAllocator {
+                allocations: Arc::clone(&allocations),
+            }));
+        let mut ctx = ExecutionCtx::new(session);
+
+        let l1 = ListArray::try_new(
+            buffer![1, 2, 3, 4].into_array(),
+            buffer![0, 3].into_array(),
+            Validity::NonNullable,
+        )
+        .unwrap();
+        let l2 = ListArray::try_new(
+            buffer![5, 6].into_array(),
+            buffer![0, 2].into_array(),
+            Validity::NonNullable,
+        )
+        .unwrap();
+
+        let chunked_list = ChunkedArray::try_new(
+            vec![l1.into_array(), l2.into_array()],
+            List(Arc::new(Primitive(I32, NonNullable)), NonNullable),
+        )
+        .unwrap()
+        .into_array();
+
+        drop(chunked_list.execute::<Canonical>(&mut ctx).unwrap());
+        assert!(
+            allocations.load(Ordering::Relaxed) >= 2,
+            "expected offset+size allocations through MemorySession"
+        );
+    }
 }

vortex-array/src/builders/mod.rs

@@ -40,6 +40,7 @@ use crate::canonical::Canonical;
 use crate::dtype::DType;
 use crate::match_each_decimal_value_type;
 use crate::match_each_native_ptype;
+use crate::memory::HostAllocatorRef;
 use crate::scalar::Scalar;
 
 mod lazy_null_builder;
@@ -291,3 +292,14 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box<dyn ArrayBui
         }
     }
 }
+
+/// Construct a new canonical builder for the given [`DType`] using a host
+/// [`crate::memory::HostAllocator`].
+pub fn builder_with_capacity_in(
+    allocator: HostAllocatorRef,
+    dtype: &DType,
+    capacity: usize,
+) -> Box<dyn ArrayBuilder> {
+    let _allocator = allocator;
+    builder_with_capacity(dtype, capacity)
+}

vortex-array/src/executor.rs

@@ -34,6 +34,8 @@ use crate::ArrayRef;
 use crate::Canonical;
 use crate::IntoArray;
 use crate::matcher::Matcher;
+use crate::memory::HostAllocatorRef;
+use crate::memory::MemorySessionExt;
 use crate::optimizer::ArrayOptimizer;
 
 /// Maximum number of iterations to attempt when executing an array before giving up and returning
@@ -210,6 +212,11 @@ impl ExecutionCtx {
         &self.session
     }
 
+    /// Get the session-scoped host allocator for this execution context.
+    pub fn allocator(&self) -> HostAllocatorRef {
+        self.session.allocator()
+    }
+
     /// Log an execution step at the current depth.
     ///
     /// Steps are accumulated and dumped as a single trace on Drop at DEBUG level.

vortex-array/src/lib.rs

@@ -52,6 +52,7 @@ pub mod kernel;
 pub mod mask;
 mod mask_future;
 pub mod matcher;
+pub mod memory;
 mod metadata;
 pub mod normalize;
 pub mod optimizer;