feat(buffer): add allocator-backed storage

Cargo.toml

@@ -97,6 +97,7 @@ rust-version = "1.95"
 version = "0.1.0"
 
 [workspace.dependencies]
+allocator-api2 = "0.2.21"
 alp = "0.0.2"
 anyhow = "1.0.100"
 arbitrary = "1.3.2"

encodings/fastlanes/src/bitpacking/array/unpack_iter.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::marker::PhantomData;
 use std::mem;
 use std::mem::MaybeUninit;
 use std::ops::Range;
@@ -51,6 +52,8 @@ impl<T: PhysicalPType<Physical: BitPacking>> UnpackStrategy<T> for BitPackingStr
 ///
 /// The usual pattern of usage should follow
 /// ```
+/// use std::mem::MaybeUninit;
+///
 /// use lending_iterator::gat;
 /// use lending_iterator::prelude::Item;
 /// #[gat(Item)]
@@ -65,17 +68,20 @@ impl<T: PhysicalPType<Physical: BitPacking>> UnpackStrategy<T> for BitPackingStr
 /// let mut ctx = vortex_array::array_session().create_execution_ctx();
 /// let array = BitPackedData::encode(&buffer![2, 3, 4, 5].into_array(), 2, &mut ctx).unwrap();
 /// let mut unpacked_chunks: BitUnpackedChunks<i32> = array.unpacked_chunks().unwrap();
+/// let mut scratch = [const { MaybeUninit::<i32>::uninit() }; 1024];
 ///
-/// if let Some(header) = unpacked_chunks.initial() {
+/// if let Some(header) = unpacked_chunks.initial(&mut scratch) {
 ///    // handle partial initial chunk
 /// }
 ///
-/// let mut chunks_iter = unpacked_chunks.full_chunks();
-/// while let Some(chunk) = chunks_iter.next() {
-///     // handle full bitpacked chunks of 1024 elements
+/// {
+///     let mut chunks_iter = unpacked_chunks.full_chunks(&mut scratch);
+///     while let Some(chunk) = chunks_iter.next() {
+///         // handle full bitpacked chunks of 1024 elements
+///     }
 /// }
 ///
-/// if let Some(trailer) = unpacked_chunks.trailer() {
+/// if let Some(trailer) = unpacked_chunks.trailer(&mut scratch) {
 ///     // handle partial trailing chunk
 /// }
 /// ```
@@ -88,7 +94,7 @@ pub struct UnpackedChunks<T: PhysicalPType, S: UnpackStrategy<T>> {
     // 0 indicates full chunk of CHUNK_SIZE
     last_chunk_length: usize,
     packed: ByteBuffer,
-    buffer: [MaybeUninit<T>; CHUNK_SIZE],
+    _marker: PhantomData<T>,
 }
 
 pub type BitUnpackedChunks<T> = UnpackedChunks<T, BitPackingStrategy>;
@@ -104,13 +110,16 @@ impl<T: BitPacked> BitUnpackedChunks<T> {
         )
     }
 
-    pub fn full_chunks(&mut self) -> BitUnpackIterator<'_, T> {
+    pub fn full_chunks<'a>(
+        &'a self,
+        scratch: &'a mut [MaybeUninit<T>; CHUNK_SIZE],
+    ) -> BitUnpackIterator<'a, T> {
         let elems_per_chunk = self.elems_per_chunk();
         let last_chunk_is_sliced = self.last_chunk_is_sliced() as usize;
         let first_chunk_is_sliced = self.first_chunk_is_sliced();
         BitUnpackIterator::new(
             buffer_as_slice(&self.packed),
-            &mut self.buffer,
+            scratch,
             self.bit_width,
             elems_per_chunk,
             self.num_chunks - last_chunk_is_sliced,
@@ -148,9 +157,9 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
             offset,
             len,
             packed,
-            buffer: [const { MaybeUninit::<T>::uninit() }; CHUNK_SIZE],
             num_chunks,
             last_chunk_length,
+            _marker: PhantomData,
         })
     }
 
@@ -160,10 +169,13 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
     }
 
     /// Access first chunk of the array if the last chunk has fewer than 1024 due to slicing
-    pub fn initial(&mut self) -> Option<&mut [T]> {
+    pub fn initial<'a>(
+        &self,
+        scratch: &'a mut [MaybeUninit<T>; CHUNK_SIZE],
+    ) -> Option<&'a mut [T]> {
         (self.first_chunk_is_sliced() || self.num_chunks == 1).then(|| {
             let chunk: &[T::Physical] = &buffer_as_slice(&self.packed)[..self.elems_per_chunk()];
-            let dst: &mut [MaybeUninit<T>] = &mut self.buffer;
+            let dst: &mut [MaybeUninit<T>] = scratch;
             let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) };
 
             let header_end_slice = if self.num_chunks == 1 {
@@ -176,17 +188,18 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
             // 2. buffer is exactly CHUNK_SIZE.
             unsafe {
                 self.strategy.unpack_chunk(self.bit_width, chunk, dst);
-                mem::transmute(&mut self.buffer[self.offset..][..header_end_slice])
+                mem::transmute(&mut scratch[self.offset..][..header_end_slice])
             }
         })
     }
 
     /// Decode all chunks (initial, full, and trailer) directly into the output range.
     pub fn decode_into(&mut self, output: &mut [MaybeUninit<T>]) {
         debug_assert_eq!(output.len(), self.len);
+        let mut scratch = [const { MaybeUninit::<T>::uninit() }; CHUNK_SIZE];
         let mut local_idx = 0;
 
-        if let Some(initial) = self.initial() {
+        if let Some(initial) = self.initial(&mut scratch) {
             local_idx = initial.len();
 
             // TODO(connor): use maybe_uninit_write_slice when it gets stabilized.
@@ -197,7 +210,7 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
 
         local_idx = self.decode_full_chunks_into_at(output, local_idx);
 
-        if let Some(trailer) = self.trailer() {
+        if let Some(trailer) = self.trailer(&mut scratch) {
             // TODO(connor): use maybe_uninit_write_slice when it gets stabilized.
             // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
             let init_trailer: &[MaybeUninit<T>] = unsafe { mem::transmute(trailer) };
@@ -226,9 +239,10 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
     where
         F: FnMut(&mut [T], Range<usize>),
     {
+        let mut scratch = [const { MaybeUninit::<T>::uninit() }; CHUNK_SIZE];
         let mut local_idx = 0;
 
-        if let Some(initial) = self.initial() {
+        if let Some(initial) = self.initial(&mut scratch) {
             let chunk_len = initial.len();
             f(initial, local_idx..local_idx + chunk_len);
             local_idx += chunk_len;
@@ -240,16 +254,16 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
             for i in self.full_chunks_range() {
                 let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk];
                 unsafe {
-                    let dst: &mut [T::Physical] = mem::transmute(&mut self.buffer[..]);
+                    let dst: &mut [T::Physical] = mem::transmute(&mut scratch[..]);
                     self.strategy.unpack_chunk(self.bit_width, chunk, dst);
-                    let unpacked: &mut [T] = mem::transmute(&mut self.buffer[..]);
+                    let unpacked: &mut [T] = mem::transmute(&mut scratch[..]);
                     f(unpacked, local_idx..local_idx + CHUNK_SIZE);
                 }
                 local_idx += CHUNK_SIZE;
             }
         }
 
-        if let Some(trailer) = self.trailer() {
+        if let Some(trailer) = self.trailer(&mut scratch) {
             let chunk_len = trailer.len();
             f(trailer, local_idx..local_idx + chunk_len);
             local_idx += chunk_len;
@@ -320,18 +334,21 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
     }
 
     /// Access last chunk of the array if the last chunk has fewer than 1024 due to slicing
-    pub fn trailer(&mut self) -> Option<&mut [T]> {
+    pub fn trailer<'a>(
+        &self,
+        scratch: &'a mut [MaybeUninit<T>; CHUNK_SIZE],
+    ) -> Option<&'a mut [T]> {
         (self.last_chunk_is_sliced() && self.num_chunks > 1).then(|| {
             let chunk: &[T::Physical] = &buffer_as_slice(&self.packed)
                 [(self.num_chunks - 1) * self.elems_per_chunk()..][..self.elems_per_chunk()];
-            let dst: &mut [MaybeUninit<T>] = &mut self.buffer;
+            let dst: &mut [MaybeUninit<T>] = scratch;
             let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) };
             // SAFETY:
             // 1. chunk is elems_per_chunk.
             // 2. buffer is exactly CHUNK_SIZE.
             unsafe {
                 self.strategy.unpack_chunk(self.bit_width, chunk, dst);
-                mem::transmute(&mut self.buffer[..self.last_chunk_length])
+                mem::transmute(&mut scratch[..self.last_chunk_length])
             }
         })
     }

encodings/fastlanes/src/bitpacking/compute/is_constant.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::mem::MaybeUninit;
 use std::ops::Range;
 
 use itertools::Itertools;
@@ -55,7 +56,8 @@ fn bitpacked_is_constant<T: BitPackedUnpack, const WIDTH: usize>(
     array: ArrayView<'_, BitPacked>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<bool> {
-    let mut bit_unpack_iterator = array.unpacked_chunks::<T>()?;
+    let bit_unpack_iterator = array.unpacked_chunks::<T>()?;
+    let mut scratch = [const { MaybeUninit::<T>::uninit() }; 1024];
     let patches = array
         .patches()
         .map(|p| -> VortexResult<_> {
@@ -68,7 +70,7 @@ fn bitpacked_is_constant<T: BitPackedUnpack, const WIDTH: usize>(
 
     let mut header_constant_value = None;
     let mut current_idx = 0;
-    if let Some(header) = bit_unpack_iterator.initial() {
+    if let Some(header) = bit_unpack_iterator.initial(&mut scratch) {
         if let Some((indices, patches, offset)) = &patches {
             apply_patches(
                 header,
@@ -87,40 +89,42 @@ fn bitpacked_is_constant<T: BitPackedUnpack, const WIDTH: usize>(
     }
 
     let mut first_chunk_value = None;
-    let mut chunks_iter = bit_unpack_iterator.full_chunks();
-    while let Some(chunk) = chunks_iter.next() {
-        if let Some((indices, patches, offset)) = &patches {
-            let chunk_len = chunk.len();
-            apply_patches(
-                chunk,
-                current_idx..current_idx + chunk_len,
-                indices,
-                patches.as_slice::<T>(),
-                *offset,
-            )
-        }
-
-        if !compute_is_constant::<_, WIDTH>(chunk) {
-            return Ok(false);
-        }
+    {
+        let mut chunks_iter = bit_unpack_iterator.full_chunks(&mut scratch);
+        while let Some(chunk) = chunks_iter.next() {
+            if let Some((indices, patches, offset)) = &patches {
+                let chunk_len = chunk.len();
+                apply_patches(
+                    chunk,
+                    current_idx..current_idx + chunk_len,
+                    indices,
+                    patches.as_slice::<T>(),
+                    *offset,
+                )
+            }
 
-        if let Some(chunk_value) = first_chunk_value {
-            if chunk_value != chunk[0] {
+            if !compute_is_constant::<_, WIDTH>(chunk) {
                 return Ok(false);
             }
-        } else {
-            if let Some(header_value) = header_constant_value
-                && header_value != chunk[0]
-            {
-                return Ok(false);
+
+            if let Some(chunk_value) = first_chunk_value {
+                if chunk_value != chunk[0] {
+                    return Ok(false);
+                }
+            } else {
+                if let Some(header_value) = header_constant_value
+                    && header_value != chunk[0]
+                {
+                    return Ok(false);
+                }
+                first_chunk_value = Some(chunk[0]);
             }
-            first_chunk_value = Some(chunk[0]);
-        }
 
-        current_idx += chunk.len();
+            current_idx += chunk.len();
+        }
     }
 
-    if let Some(trailer) = bit_unpack_iterator.trailer() {
+    if let Some(trailer) = bit_unpack_iterator.trailer(&mut scratch) {
         if let Some((indices, patches, offset)) = &patches {
             let chunk_len = trailer.len();
             apply_patches(

encodings/pco/src/array.rs

@@ -47,7 +47,6 @@ use vortex_array::vtable::child_to_validity;
 use vortex_array::vtable::validity_to_child;
 use vortex_buffer::BufferMut;
 use vortex_buffer::ByteBuffer;
-use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexError;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -568,17 +567,17 @@ impl PcoData {
                 }
             );
 
-            let mut chunk_meta_buffer = ByteBufferMut::with_capacity(cc.meta_size_hint());
+            let mut chunk_meta_buffer = Vec::with_capacity(cc.meta_size_hint());
             cc.write_meta(&mut chunk_meta_buffer)
                 .map_err(vortex_err_from_pco)?;
-            chunk_meta_buffers.push(chunk_meta_buffer.freeze());
+            chunk_meta_buffers.push(ByteBuffer::from(chunk_meta_buffer));
 
             let mut page_infos = vec![];
             for (page_idx, page_n_values) in cc.n_per_page().into_iter().enumerate() {
-                let mut page = ByteBufferMut::with_capacity(cc.page_size_hint(page_idx));
+                let mut page = Vec::with_capacity(cc.page_size_hint(page_idx));
                 cc.write_page(page_idx, &mut page)
                     .map_err(vortex_err_from_pco)?;
-                page_buffers.push(page.freeze());
+                page_buffers.push(ByteBuffer::from(page));
                 page_infos.push(PcoPageInfo {
                     n_values: u32::try_from(page_n_values)?,
                 });

encodings/sparse/src/lib.rs

@@ -48,7 +48,6 @@ use vortex_array::validity::Validity;
 use vortex_array::vtable::VTable;
 use vortex_array::vtable::ValidityVTable;
 use vortex_buffer::Buffer;
-use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexExpect as _;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -217,7 +216,7 @@ impl VTable for Sparse {
         match idx {
             0 => {
                 let fill_value_buffer =
-                    ScalarValue::to_proto_bytes::<ByteBufferMut>(array.fill_value.value()).freeze();
+                    ScalarValue::to_proto_bytes::<Vec<u8>>(array.fill_value.value()).into();
                 BufferHandle::new_host(fill_value_buffer)
             }
             _ => vortex_panic!("SparseArray buffer index {idx} out of bounds"),

encodings/zstd/src/array.rs

@@ -1144,7 +1144,7 @@ impl ZstdData {
 
         let value_bytes = values.buffer_handle().try_to_host_sync()?;
         // Align frames to buffer alignment. This is necessary for overaligned buffers.
-        let alignment = *value_bytes.alignment();
+        let alignment = value_bytes.alignment().as_usize();
         let step_width = (values_per_frame * byte_width).div_ceil(alignment) * alignment;
 
         let frame_byte_starts = (0..n_values * byte_width)

encodings/zstd/src/zstd_buffers.rs

@@ -357,7 +357,7 @@ fn compute_output_layout(
     let mut total_size = 0usize;
 
     for (&size, &alignment) in output_sizes.iter().zip(output_alignments.iter()) {
-        total_size = total_size.next_multiple_of(*alignment);
+        total_size = total_size.next_multiple_of(alignment.as_usize());
         offsets.push(total_size);
         total_size += size;
     }

fuzz/fuzz_targets/file_io.rs

@@ -19,7 +19,6 @@ use vortex_array::expr::lit;
 use vortex_array::expr::root;
 use vortex_array::scalar_fn::fns::operators::Operator;
 use vortex_btrblocks::BtrBlocksCompressorBuilder;
-use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::vortex_panic;
 use vortex_file::OpenOptionsSessionExt;
@@ -72,7 +71,7 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus {
         ),
     };
 
-    let mut full_buff = ByteBufferMut::empty();
+    let mut full_buff = Vec::new();
     let _footer = write_options
         .blocking(&*RUNTIME)
         .write(&mut full_buff, array_data.to_array_iterator())

vortex-array/src/arrays/constant/vtable/mod.rs

@@ -6,7 +6,6 @@ use std::hash::Hash;
 use std::hash::Hasher;
 
 use itertools::Itertools;
-use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
@@ -107,7 +106,7 @@ impl VTable for Constant {
     fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
         match idx {
             0 => BufferHandle::new_host(
-                ScalarValue::to_proto_bytes::<ByteBufferMut>(array.scalar.value()).freeze(),
+                ScalarValue::to_proto_bytes::<Vec<u8>>(array.scalar.value()).into(),
             ),
             _ => vortex_panic!("ConstantArray buffer index {idx} out of bounds"),
         }