refactor(array): use buffer allocator in execution context

vortex-array/Cargo.toml

@@ -83,6 +83,7 @@ serde = ["dep:serde", "vortex-buffer/serde", "vortex-mask/serde"]
 unstable_row_fns = []
 
 [dev-dependencies]
+allocator-api2 = { workspace = true }
 divan = { workspace = true }
 futures = { workspace = true, features = ["executor"] }
 insta = { workspace = true }

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

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use itertools::Itertools as _;
-use vortex_buffer::Buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
@@ -30,7 +29,6 @@ use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::Nullability;
 use crate::dtype::PType;
-use crate::memory::HostAllocatorExt;
 use crate::validity::Validity;
 
 pub(super) fn _canonicalize(
@@ -71,7 +69,8 @@ pub(super) fn _canonicalize(
         }
         DType::Variant(_) => Canonical::Variant(pack_variant_chunks(owned_chunks, ctx)?),
         _ => {
-            let mut builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len());
+            let mut builder =
+                builder_with_capacity_in(ctx.allocator().clone(), array.dtype(), array.len());
             array.array().append_to_builder(builder.as_mut(), ctx)?;
             builder.finish_into_canonical(ctx)
         }
@@ -180,10 +179,10 @@ fn swizzle_list_chunks(
     // We (somewhat arbitrarily) choose `u64` for our offsets and sizes here. These can always be
     // narrowed later by the compressor.
     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 offsets = allocator.zeroed::<u64>(len);
+    let mut sizes = allocator.zeroed::<u64>(len);
+    let offsets_out = offsets.as_mut_slice();
+    let sizes_slice_out = sizes.as_mut_slice();
     let mut next_list = 0usize;
 
     for chunk in chunks {
@@ -229,16 +228,8 @@ fn swizzle_list_chunks(
         unsafe { ChunkedArray::new_unchecked(list_elements_chunks, elem_dtype.clone()) }
             .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();
+    let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable).into_array();
+    let sizes = PrimitiveArray::new(sizes.freeze(), Validity::NonNullable).into_array();
 
     // SAFETY:
     // - `offsets` and `sizes` are non-nullable u64 arrays of the same length
@@ -286,11 +277,17 @@ fn swizzle_fixed_size_list_chunks(
 
 #[cfg(test)]
 mod tests {
+    use std::alloc::Layout;
+    use std::ptr::NonNull;
     use std::sync::Arc;
     use std::sync::LazyLock;
     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 vortex_buffer::BufferAllocatorRef;
     use vortex_buffer::buffer;
     use vortex_error::VortexResult;
     use vortex_error::vortex_bail;
@@ -318,10 +315,7 @@ mod tests {
     use crate::dtype::DType::Variant as VariantDType;
     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::scalar::Scalar;
     use crate::validity::Validity;
 
@@ -333,14 +327,16 @@ mod tests {
         allocations: Arc<AtomicUsize>,
     }
 
-    impl HostAllocator for CountingAllocator {
-        fn allocate(
-            &self,
-            len: usize,
-            alignment: vortex_buffer::Alignment,
-        ) -> VortexResult<WritableHostBuffer> {
+    // SAFETY: this forwards memory operations to Global and only counts allocations.
+    unsafe impl Allocator for CountingAllocator {
+        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
             self.allocations.fetch_add(1, Ordering::Relaxed);
-            DefaultHostAllocator.allocate(len, alignment)
+            Global.allocate(layout)
+        }
+
+        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+            // SAFETY: ptr and layout came from Global.
+            unsafe { Global.deallocate(ptr, layout) }
         }
     }
 
@@ -663,9 +659,10 @@ mod tests {
     #[test]
     fn list_canonicalize_uses_memory_session_allocator() {
         let allocations = Arc::new(AtomicUsize::new(0));
-        let session = crate::array_session().with_allocator(Arc::new(CountingAllocator {
-            allocations: Arc::clone(&allocations),
-        }));
+        let session =
+            crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator {
+                allocations: Arc::clone(&allocations),
+            }));
         let mut ctx = session.create_execution_ctx();
 
         let l1 = ListArray::try_new(

vortex-array/src/builders/mod.rs

@@ -47,7 +47,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::memory::BufferAllocatorRef;
 use crate::scalar::Scalar;
 
 mod lazy_null_builder;
@@ -447,9 +447,9 @@ 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`].
+/// [`vortex_buffer::BufferAllocator`].
 pub fn builder_with_capacity_in(
-    allocator: HostAllocatorRef,
+    allocator: BufferAllocatorRef,
     dtype: &DType,
     capacity: usize,
 ) -> Box<dyn ArrayBuilder> {

vortex-array/src/executor.rs

@@ -17,6 +17,7 @@ use std::fmt;
 use std::fmt::Display;
 use std::sync::Arc;
 use std::sync::LazyLock;
+use std::sync::OnceLock;
 #[cfg(debug_assertions)]
 use std::sync::atomic::AtomicUsize;
 #[cfg(debug_assertions)]
@@ -38,7 +39,7 @@ use crate::builders::ArrayBuilder;
 use crate::builders::builder_with_capacity_in;
 use crate::dtype::DType;
 use crate::matcher::Matcher;
-use crate::memory::HostAllocatorRef;
+use crate::memory::BufferAllocatorRef;
 use crate::memory::MemorySessionExt;
 use crate::optimizer::ArrayOptimizer;
 use crate::optimizer::kernels::ArrayKernelsExt;
@@ -292,7 +293,7 @@ impl ArrayRef {
                     if current_builder.is_none() {
                         trace_op!(record_builder_start(&array));
                         current_builder = Some(builder_with_capacity_in(
-                            ctx.allocator(),
+                            ctx.allocator().clone(),
                             array.dtype(),
                             array.len(),
                         ));
@@ -350,6 +351,8 @@ struct StackFrame {
 #[derive(Debug, Clone)]
 pub struct ExecutionCtx {
     session: VortexSession,
+    // OnceLock avoids cloning the session allocator when a context does not allocate.
+    allocator: OnceLock<BufferAllocatorRef>,
     execute_parent_kernels: Arc<ParentExecutionKernels>,
     #[cfg(debug_assertions)]
     id: usize,
@@ -367,6 +370,7 @@ impl ExecutionCtx {
         let execute_parent_kernels = session.kernels().execute_parent_snapshot();
         Self {
             session,
+            allocator: OnceLock::new(),
             execute_parent_kernels,
             #[cfg(debug_assertions)]
             id: {
@@ -383,9 +387,15 @@ impl ExecutionCtx {
         &self.session
     }
 
-    /// Get the session-scoped host allocator for this execution context.
-    pub fn allocator(&self) -> HostAllocatorRef {
-        self.session.allocator()
+    /// Get the allocator for this execution context.
+    pub fn allocator(&self) -> &BufferAllocatorRef {
+        self.allocator.get_or_init(|| self.session.allocator())
+    }
+
+    /// Set the allocator for this execution context.
+    pub fn with_allocator(mut self, allocator: BufferAllocatorRef) -> Self {
+        self.allocator = OnceLock::from(allocator);
+        self
     }
 
     /// Log an execution step at the current depth.
@@ -531,7 +541,8 @@ impl Executable for ArrayRef {
             ExecutionStep::AppendChild(_) => {
                 // Single-step: build the entire parent via the builder path.
                 trace_op!(record_builder_start(&array));
-                let builder = builder_with_capacity_in(ctx.allocator(), array.dtype(), array.len());
+                let builder =
+                    builder_with_capacity_in(ctx.allocator().clone(), array.dtype(), array.len());
                 let mut builder = execute_into_builder(array, builder, ctx)?;
                 let output = builder.finish();
                 trace_op!(record_builder_finish(&output));
@@ -926,17 +937,24 @@ impl VortexSessionExecute for VortexSession {
 
 #[cfg(test)]
 mod tests {
+    use static_assertions::assert_impl_all;
+    use vortex_session::SessionExt;
     use vortex_session::VortexSession;
 
     use super::*;
     use crate::VTable as _;
     use crate::VortexSessionExecute;
     use crate::arrays::Bool;
     use crate::arrays::Primitive;
+    use crate::memory::BufferAllocatorRef;
+    use crate::memory::MemorySession;
+    use crate::memory::MemorySessionExt;
     use crate::optimizer::kernels::ExecuteParentFn;
     use crate::optimizer::kernels::KernelSession;
     use crate::optimizer::kernels::execute_parent_key;
 
+    assert_impl_all!(ExecutionCtx: Send, Sync);
+
     fn noop_execute_parent(
         _child: &ArrayRef,
         _parent: &ArrayRef,
@@ -974,4 +992,27 @@ mod tests {
         let after_registration = session.create_execution_ctx();
         assert!(after_registration.execute_parent_kernels.contains_key(&key));
     }
+
+    #[test]
+    fn execution_ctx_allocator_override() {
+        let first = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator);
+        let second = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator);
+        let third = BufferAllocatorRef::new(vortex_buffer::StaticBufferAllocator);
+        let session = VortexSession::empty()
+            .with::<MemorySession>()
+            .with_allocator(first.clone());
+        let ctx = session.create_execution_ctx();
+
+        session
+            .get_mut::<MemorySession>()
+            .set_allocator(third.clone());
+
+        assert!(session.allocator().ptr_eq(&third));
+        assert!(ctx.allocator().ptr_eq(&third));
+
+        let ctx = ctx.with_allocator(second.clone());
+        session.get_mut::<MemorySession>().set_allocator(first);
+
+        assert!(ctx.allocator().ptr_eq(&second));
+    }
 }