Initial buffer allocator

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::BufferAllocatorRef, 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::DefaultBufferAllocator
+
+impl core::default::Default for vortex_array::memory::DefaultBufferAllocator
+
+pub fn vortex_array::memory::DefaultBufferAllocator::default() -> vortex_array::memory::DefaultBufferAllocator
+
+impl core::fmt::Debug for vortex_array::memory::DefaultBufferAllocator
+
+pub fn vortex_array::memory::DefaultBufferAllocator::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::memory::BufferAllocator for vortex_array::memory::DefaultBufferAllocator
+
+pub fn vortex_array::memory::DefaultBufferAllocator::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::BufferAllocatorRef
+
+pub fn vortex_array::memory::MemorySession::new(allocator: vortex_array::memory::BufferAllocatorRef) -> Self
+
+pub fn vortex_array::memory::MemorySession::set_allocator(&mut self, allocator: vortex_array::memory::BufferAllocatorRef)
+
+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::BufferAllocator: core::fmt::Debug + core::marker::Send + core::marker::Sync + 'static
+
+pub fn vortex_array::memory::BufferAllocator::allocate(&self, len: usize, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+impl vortex_array::memory::BufferAllocator for vortex_array::memory::DefaultBufferAllocator
+
+pub fn vortex_array::memory::DefaultBufferAllocator::allocate(&self, len: usize, alignment: vortex_buffer::alignment::Alignment) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+pub trait vortex_array::memory::BufferAllocatorExt: vortex_array::memory::BufferAllocator
+
+pub fn vortex_array::memory::BufferAllocatorExt::allocate_typed<T>(&self, len: usize) -> vortex_error::VortexResult<vortex_array::memory::WritableHostBuffer>
+
+impl<A: vortex_array::memory::BufferAllocator + ?core::marker::Sized> vortex_array::memory::BufferAllocatorExt 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::BufferAllocatorRef
+
+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::BufferAllocatorRef
+
+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::BufferAllocatorRef = alloc::sync::Arc<dyn vortex_array::memory::BufferAllocator>
+
 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::BufferAllocatorRef
+
 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/memory.rs

@@ -10,6 +10,7 @@ use std::sync::Arc;
 
 use bytes::Bytes;
 use vortex_buffer::Alignment;
+use vortex_buffer::Buffer;
 use vortex_buffer::ByteBuffer;
 use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexResult;
@@ -24,6 +25,11 @@ pub trait HostBufferMut: Send + 'static {
     /// Returns the logical byte length of the buffer.
     fn len(&self) -> usize;
 
+    /// Whether the buffer is empty.
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
     /// Returns the alignment of the buffer.
     fn alignment(&self) -> Alignment;
 
@@ -72,25 +78,23 @@ impl WritableHostBuffer {
             InvalidArgument: "Cannot create typed mutable slice for zero-sized type {}",
             std::any::type_name::<T>()
         );
+        vortex_ensure!(
+            self.alignment().is_aligned_to(Alignment::of::<T>()),
+            InvalidArgument: "Buffer is not sufficiently aligned for type {}",
+            std::any::type_name::<T>()
+        );
 
         let bytes = self.as_mut_slice();
         let byte_len = bytes.len();
         let ptr = bytes.as_mut_ptr();
         let type_size = size_of::<T>();
-        let type_align = align_of::<T>();
 
         vortex_ensure!(
-            byte_len % type_size == 0,
+            byte_len.is_multiple_of(type_size),
             InvalidArgument: "Buffer length {byte_len} is not a multiple of {} for {}",
             type_size,
             std::any::type_name::<T>()
         );
-        vortex_ensure!(
-            (ptr as usize) % type_align == 0,
-            InvalidArgument: "Buffer pointer is not aligned to {} for {}",
-            type_align,
-            std::any::type_name::<T>()
-        );
 
         // SAFETY: We checked size divisibility and pointer alignment for `T`,
         // and we have exclusive mutable access to the underlying bytes.
@@ -101,6 +105,35 @@ impl WritableHostBuffer {
     pub fn freeze(self) -> ByteBuffer {
         self.inner.freeze()
     }
+
+    /// Freeze the writable buffer into a typed immutable [`Buffer<T>`].
+    pub fn freeze_typed<T>(self) -> VortexResult<Buffer<T>> {
+        vortex_ensure!(
+            size_of::<T>() != 0,
+            InvalidArgument: "Cannot freeze typed buffer for zero-sized type {}",
+            std::any::type_name::<T>()
+        );
+
+        let buffer = self.freeze();
+        let byte_len = buffer.len();
+        let type_size = size_of::<T>();
+        let type_align = Alignment::of::<T>();
+
+        vortex_ensure!(
+            byte_len % type_size == 0,
+            InvalidArgument: "Buffer length {byte_len} is not a multiple of {} for {}",
+            type_size,
+            std::any::type_name::<T>()
+        );
+        vortex_ensure!(
+            buffer.is_aligned(type_align),
+            InvalidArgument: "Buffer pointer is not aligned to {} for {}",
+            type_align,
+            std::any::type_name::<T>()
+        );
+
+        Ok(Buffer::from_byte_buffer(buffer))
+    }
 }
 
 impl Debug for WritableHostBuffer {
@@ -320,4 +353,26 @@ mod tests {
         let msg = format!("{err}");
         assert!(msg.contains("not a multiple of"));
     }
+
+    #[test]
+    fn freeze_typed_round_trip() {
+        let allocator = DefaultBufferAllocator;
+        let mut writable = allocator.allocate_typed::<u64>(4).unwrap();
+        writable
+            .as_mut_slice_typed::<u64>()
+            .unwrap()
+            .copy_from_slice(&[1, 3, 5, 7]);
+
+        let frozen = writable.freeze_typed::<u64>().unwrap();
+        assert_eq!(frozen.as_slice(), [1, 3, 5, 7]);
+    }
+
+    #[test]
+    fn freeze_typed_rejects_length_mismatch() {
+        let allocator = DefaultBufferAllocator;
+        let writable = allocator.allocate(7, Alignment::none()).unwrap();
+        let err = writable.freeze_typed::<u32>().unwrap_err();
+        let msg = format!("{err}");
+        assert!(msg.contains("not a multiple of"));
+    }
 }

vortex-file/public-api.lock

@@ -398,11 +398,11 @@ pub const vortex_file::VORTEX_FILE_EXTENSION: &str
 
 pub static vortex_file::ALLOWED_ENCODINGS: std::sync::lazy_lock::LazyLock<vortex_utils::aliases::hash_set::HashSet<vortex_array::array::ArrayId>>
 
-pub trait vortex_file::OpenOptionsSessionExt: vortex_array::session::ArraySessionExt + vortex_layout::session::LayoutSessionExt + vortex_io::session::RuntimeSessionExt
+pub trait vortex_file::OpenOptionsSessionExt: vortex_array::session::ArraySessionExt + vortex_layout::session::LayoutSessionExt + vortex_io::session::RuntimeSessionExt + vortex_array::memory::MemorySessionExt
 
 pub fn vortex_file::OpenOptionsSessionExt::open_options(&self) -> vortex_file::VortexOpenOptions
 
-impl<S: vortex_array::session::ArraySessionExt + vortex_layout::session::LayoutSessionExt + vortex_io::session::RuntimeSessionExt> vortex_file::OpenOptionsSessionExt for S
+impl<S: vortex_array::session::ArraySessionExt + vortex_layout::session::LayoutSessionExt + vortex_io::session::RuntimeSessionExt + vortex_array::memory::MemorySessionExt> vortex_file::OpenOptionsSessionExt for S
 
 pub fn S::open_options(&self) -> vortex_file::VortexOpenOptions
 

vortex-io/public-api.lock

@@ -150,6 +150,8 @@ impl vortex_io::object_store::ObjectStoreReadAt
 
 pub fn vortex_io::object_store::ObjectStoreReadAt::new(store: alloc::sync::Arc<dyn object_store::ObjectStore>, path: object_store::path::Path, handle: vortex_io::runtime::Handle) -> Self
 
+pub fn vortex_io::object_store::ObjectStoreReadAt::new_with_allocator(store: alloc::sync::Arc<dyn object_store::ObjectStore>, path: object_store::path::Path, handle: vortex_io::runtime::Handle, allocator: vortex_array::memory::BufferAllocatorRef) -> Self
+
 pub fn vortex_io::object_store::ObjectStoreReadAt::with_coalesce_config(self, config: vortex_io::CoalesceConfig) -> Self
 
 pub fn vortex_io::object_store::ObjectStoreReadAt::with_concurrency(self, concurrency: usize) -> Self
@@ -466,6 +468,8 @@ impl vortex_io::std_file::FileReadAt
 
 pub fn vortex_io::std_file::FileReadAt::open(path: impl core::convert::AsRef<std::path::Path>, handle: vortex_io::runtime::Handle) -> vortex_error::VortexResult<Self>
 
+pub fn vortex_io::std_file::FileReadAt::open_with_allocator(path: impl core::convert::AsRef<std::path::Path>, handle: vortex_io::runtime::Handle, allocator: vortex_array::memory::BufferAllocatorRef) -> vortex_error::VortexResult<Self>
+
 impl vortex_io::VortexReadAt for vortex_io::std_file::FileReadAt
 
 pub fn vortex_io::std_file::FileReadAt::coalesce_config(&self) -> core::option::Option<vortex_io::CoalesceConfig>