fix(buffer): address allocator review follow-ups

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

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use itertools::Itertools as _;
+use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
@@ -179,8 +180,8 @@ 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.zeroed::<u64>(len);
-    let mut sizes = allocator.zeroed::<u64>(len);
+    let mut offsets = BufferMut::<u64>::zeroed_in(len, allocator.clone());
+    let mut sizes = BufferMut::<u64>::zeroed_in(len, allocator.clone());
     let offsets_out = offsets.as_mut_slice();
     let sizes_slice_out = sizes.as_mut_slice();
     let mut next_list = 0usize;
@@ -659,10 +660,11 @@ mod tests {
     #[test]
     fn list_canonicalize_uses_memory_session_allocator() {
         let allocations = Arc::new(AtomicUsize::new(0));
-        let session =
-            crate::array_session().with_allocator(BufferAllocatorRef::new(CountingAllocator {
+        let session = crate::array_session().with_allocator(BufferAllocatorRef::new_arc(Arc::new(
+            CountingAllocator {
                 allocations: Arc::clone(&allocations),
-            }));
+            },
+        )));
         let mut ctx = session.create_execution_ctx();
 
         let l1 = ListArray::try_new(

vortex-array/src/executor.rs

@@ -995,9 +995,9 @@ mod tests {
 
     #[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 first = BufferAllocatorRef::new_arc(Arc::new(vortex_buffer::StaticBufferAllocator));
+        let second = BufferAllocatorRef::new_arc(Arc::new(vortex_buffer::StaticBufferAllocator));
+        let third = BufferAllocatorRef::new_arc(Arc::new(vortex_buffer::StaticBufferAllocator));
         let session = VortexSession::empty()
             .with::<MemorySession>()
             .with_allocator(first.clone());
@@ -1007,12 +1007,12 @@ mod tests {
             .get_mut::<MemorySession>()
             .set_allocator(third.clone());
 
-        assert!(session.allocator().ptr_eq(&third));
-        assert!(ctx.allocator().ptr_eq(&third));
+        assert!(std::ptr::eq(session.allocator().as_ref(), third.as_ref()));
+        assert!(std::ptr::eq(ctx.allocator().as_ref(), third.as_ref()));
 
         let ctx = ctx.with_allocator(second.clone());
         session.get_mut::<MemorySession>().set_allocator(first);
 
-        assert!(ctx.allocator().ptr_eq(&second));
+        assert!(std::ptr::eq(ctx.allocator().as_ref(), second.as_ref()));
     }
 }

vortex-array/src/memory.rs

@@ -7,6 +7,7 @@ use std::any::Any;
 
 pub use vortex_buffer::BufferAllocator;
 pub use vortex_buffer::BufferAllocatorRef;
+pub use vortex_buffer::DEFAULT_BUFFER_ALLOCATOR;
 pub use vortex_buffer::StaticBufferAllocator;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
@@ -38,7 +39,7 @@ impl MemorySession {
 
 impl Default for MemorySession {
     fn default() -> Self {
-        Self::new(BufferAllocatorRef::statically_allocated())
+        Self::new(DEFAULT_BUFFER_ALLOCATOR.clone())
     }
 }
 
@@ -76,16 +77,17 @@ impl<S: SessionExt> MemorySessionExt for S {}
 
 #[cfg(test)]
 mod tests {
-    use vortex_buffer::BufferAllocatorRef;
+    use vortex_buffer::BufferMut;
+    use vortex_buffer::DEFAULT_BUFFER_ALLOCATOR;
 
     use super::MemorySession;
 
     #[test]
     fn memory_session_replaces_allocator() {
-        let allocator = BufferAllocatorRef::statically_allocated();
+        let allocator = DEFAULT_BUFFER_ALLOCATOR.clone();
         let mut session = MemorySession::default();
         session.set_allocator(allocator);
-        let buffer = session.allocator().copy_from([1u32, 2, 3]);
+        let buffer = BufferMut::copy_from_in([1u32, 2, 3], session.allocator());
         assert_eq!(buffer.as_slice(), [1, 2, 3]);
     }
 }

vortex-buffer/Cargo.toml

@@ -24,6 +24,7 @@ warn-copy = ["dep:tracing"]
 
 [dependencies]
 allocator-api2 = { workspace = true }
+arcref = { workspace = true }
 arrow-buffer = { workspace = true }
 bitvec = { workspace = true }
 bytes = { workspace = true }

vortex-buffer/benches/allocation.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
 use allocator_api2::alloc::Global;
 use arrow_buffer::MutableBuffer;
 use bytes::BytesMut;
@@ -9,8 +11,10 @@ use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
+use vortex_buffer::StaticBufferAllocator;
 
 const SIZES: &[usize] = &[0, 64, 256, 1024, 16_384, 65_536];
+static GLOBAL_ALLOCATOR: Global = Global;
 
 fn main() {
     divan::main();
@@ -22,19 +26,20 @@ fn allocate_drop_vortex(bencher: Bencher, size: usize) {
 }
 
 #[divan::bench(args = SIZES)]
-fn allocate_drop_vortex_custom(bencher: Bencher, size: usize) {
+fn allocate_drop_vortex_arc(bencher: Bencher, size: usize) {
     bencher
-        .with_inputs(|| BufferAllocatorRef::new(Global))
-        .bench_refs(|allocator| drop(allocator.with_capacity::<u8>(size)));
+        .with_inputs(|| BufferAllocatorRef::new_arc(Arc::new(StaticBufferAllocator)))
+        .bench_refs(|allocator| drop(BufferMut::<u8>::with_capacity_in(size, allocator.clone())));
 }
 
 #[divan::bench(args = SIZES)]
 fn allocate_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) {
     bencher.bench(|| {
-        drop(BufferMut::<u8>::with_capacity_preferred_aligned(
+        drop(BufferMut::<u8>::with_capacity_preferred_aligned_in(
             size,
             Alignment::of::<u8>(),
             None,
+            BufferAllocatorRef::new_ref(&GLOBAL_ALLOCATOR),
         ))
     });
 }
@@ -55,18 +60,25 @@ fn allocate_freeze_drop_vortex(bencher: Bencher, size: usize) {
 }
 
 #[divan::bench(args = SIZES)]
-fn allocate_freeze_drop_vortex_custom(bencher: Bencher, size: usize) {
+fn allocate_freeze_drop_vortex_arc(bencher: Bencher, size: usize) {
     bencher
-        .with_inputs(|| BufferAllocatorRef::new(Global))
-        .bench_refs(|allocator| drop(allocator.with_capacity::<u8>(size).freeze()));
+        .with_inputs(|| BufferAllocatorRef::new_arc(Arc::new(StaticBufferAllocator)))
+        .bench_refs(|allocator| {
+            drop(BufferMut::<u8>::with_capacity_in(size, allocator.clone()).freeze())
+        });
 }
 
 #[divan::bench(args = SIZES)]
 fn allocate_freeze_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) {
     bencher.bench(|| {
         drop(
-            BufferMut::<u8>::with_capacity_preferred_aligned(size, Alignment::of::<u8>(), None)
-                .freeze(),
+            BufferMut::<u8>::with_capacity_preferred_aligned_in(
+                size,
+                Alignment::of::<u8>(),
+                None,
+                BufferAllocatorRef::new_ref(&GLOBAL_ALLOCATOR),
+            )
+            .freeze(),
         )
     });
 }

vortex-buffer/src/alignment.rs

@@ -25,7 +25,7 @@ impl Alignment {
     /// Default alignment for device-to-host buffer copies.
     pub const HOST_COPY: Self = Alignment::new(256);
 
-    /// Default alignment for all buffers.
+    /// Default preferred alignment for Vortex buffers.
     ///
     /// Chosen to be larger than any SIMD register (e.g. AVX-512's 64-byte
     /// registers) so that buffers can be processed with vectorized loads/stores