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/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"),

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"),
         }

vortex-array/src/arrays/struct_/vtable/operations.rs

@@ -9,19 +9,25 @@ use crate::array::OperationsVTable;
 use crate::arrays::Struct;
 use crate::arrays::struct_::StructArrayExt;
 use crate::scalar::Scalar;
+use crate::scalar::ScalarValue;
 
 impl OperationsVTable<Struct> for Struct {
     fn scalar_at(
         array: ArrayView<'_, Struct>,
         index: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar> {
-        let field_scalars: VortexResult<Vec<Scalar>> = array
+        let field_values = array
             .iter_unmasked_fields()
-            .map(|field| field.execute_scalar(index, ctx))
-            .collect();
+            .map(|field| field.execute_scalar(index, ctx).map(Scalar::into_value))
+            .collect::<VortexResult<Vec<_>>>()?;
         // SAFETY: The vtable guarantees index is in-bounds and non-null before this is called.
-        // Each field's scalar_at returns a scalar with the field's own dtype.
-        Ok(unsafe { Scalar::struct_unchecked(array.dtype().clone(), field_scalars?) })
+        // Each field's scalar_at returns a value with the field's own dtype.
+        Ok(unsafe {
+            Scalar::new_unchecked(
+                array.dtype().clone(),
+                Some(ScalarValue::Tuple(field_values)),
+            )
+        })
     }
 }

vortex-array/src/builders/varbinview.rs

@@ -800,7 +800,8 @@ impl ArrayBuilder for VarBinViewBuilder {
 }
 
 impl VarBinViewBuilder {
-    #[inline]
+    #[allow(clippy::inline_always)]
+    #[inline(always)]
     fn push_view(
         &mut self,
         view: BinaryView,

vortex-array/src/memory.rs

@@ -8,7 +8,6 @@ use std::fmt::Debug;
 use std::mem::size_of;
 use std::sync::Arc;
 
-use bytes::Bytes;
 use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
 use vortex_buffer::ByteBuffer;
@@ -244,26 +243,14 @@ impl HostAllocator for DefaultHostAllocator {
         // SAFETY: We fully initialize this slice before freezing it.
         unsafe { buffer.set_len(len) };
         Ok(WritableHostBuffer::new(Box::new(
-            DefaultWritableHostBuffer { buffer, alignment },
+            DefaultWritableHostBuffer { buffer },
         )))
     }
 }
 
 #[derive(Debug)]
 struct DefaultWritableHostBuffer {
     buffer: ByteBufferMut,
-    alignment: Alignment,
-}
-
-#[derive(Debug)]
-struct HostBufferOwner {
-    buffer: ByteBufferMut,
-}
-
-impl AsRef<[u8]> for HostBufferOwner {
-    fn as_ref(&self) -> &[u8] {
-        self.buffer.as_slice()
-    }
 }
 
 impl HostBufferMut for DefaultWritableHostBuffer {
@@ -272,17 +259,15 @@ impl HostBufferMut for DefaultWritableHostBuffer {
     }
 
     fn alignment(&self) -> Alignment {
-        self.alignment
+        self.buffer.alignment()
     }
 
     fn as_mut_slice(&mut self) -> &mut [u8] {
         self.buffer.as_mut_slice()
     }
 
     fn freeze(self: Box<Self>) -> ByteBuffer {
-        let Self { buffer, alignment } = *self;
-        let bytes = Bytes::from_owner(HostBufferOwner { buffer });
-        ByteBuffer::from_bytes_aligned(bytes, alignment)
+        self.buffer.freeze()
     }
 }
 

vortex-arrow/src/executor/byte_view.rs

@@ -6,12 +6,12 @@ use std::sync::Arc;
 use arrow_array::ArrayRef as ArrowArrayRef;
 use arrow_array::GenericByteViewArray;
 use arrow_array::types::ByteViewType;
-use arrow_buffer::ScalarBuffer;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::dtype::Nullability;
+use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 
 use crate::dtype::from_arrow_data_type;
@@ -22,8 +22,8 @@ pub fn canonical_varbinview_to_arrow<T: ByteViewType>(
     array: &VarBinViewArray,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrowArrayRef> {
-    let views =
-        ScalarBuffer::<u128>::from(array.views_handle().as_host().clone().into_arrow_buffer());
+    let views = Buffer::<u128>::from_byte_buffer(array.views_handle().as_host().clone())
+        .into_arrow_scalar_buffer();
     let buffers: Vec<_> = array
         .data_buffers()
         .iter()
@@ -64,3 +64,23 @@ pub(super) fn to_arrow_byte_view<T: ByteViewType>(
     let varbinview = array.execute::<VarBinViewArray>(ctx)?;
     execute_varbinview_to_arrow::<T>(&varbinview, ctx)
 }
+
+#[cfg(test)]
+mod tests {
+    use arrow_array::types::StringViewType;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::array_session;
+
+    use super::*;
+
+    #[test]
+    fn empty_views_are_aligned() -> VortexResult<()> {
+        let array = VarBinViewArray::from_iter_str(std::iter::empty::<&str>());
+        let mut ctx = array_session().create_execution_ctx();
+
+        let arrow = canonical_varbinview_to_arrow::<StringViewType>(&array, &mut ctx)?;
+
+        assert!(arrow.is_empty());
+        Ok(())
+    }
+}

vortex-buffer/Cargo.toml

@@ -23,6 +23,7 @@ serde = ["dep:serde", "serde/serde_derive"]
 warn-copy = ["dep:tracing"]
 
 [dependencies]
+allocator-api2 = { workspace = true }
 arrow-buffer = { workspace = true }
 bitvec = { workspace = true }
 bytes = { workspace = true }
@@ -57,3 +58,7 @@ harness = false
 [[bench]]
 name = "collect_bool"
 harness = false
+
+[[bench]]
+name = "allocation"
+harness = false

vortex-buffer/benches/allocation.rs

@@ -0,0 +1,106 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use allocator_api2::alloc::Global;
+use arrow_buffer::MutableBuffer;
+use bytes::BytesMut;
+use divan::Bencher;
+use vortex_buffer::Alignment;
+use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
+use vortex_buffer::BufferMut;
+
+const SIZES: &[usize] = &[0, 64, 256, 1024, 16_384, 65_536];
+
+fn main() {
+    divan::main();
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_drop_vortex(bencher: Bencher, size: usize) {
+    bencher.bench(|| drop(BufferMut::<u8>::with_capacity(size)));
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_drop_vortex_custom(bencher: Bencher, size: usize) {
+    bencher
+        .with_inputs(|| BufferAllocatorRef::new(Global))
+        .bench_refs(|allocator| drop(allocator.with_capacity::<u8>(size)));
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_drop_vortex_minimal_alignment(bencher: Bencher, size: usize) {
+    bencher.bench(|| {
+        drop(BufferMut::<u8>::with_capacity_preferred_aligned(
+            size,
+            Alignment::of::<u8>(),
+            None,
+        ))
+    });
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_drop_bytes(bencher: Bencher, size: usize) {
+    bencher.bench(|| drop(BytesMut::with_capacity(size)));
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_drop_arrow(bencher: Bencher, size: usize) {
+    bencher.bench(|| drop(MutableBuffer::with_capacity(size)));
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_freeze_drop_vortex(bencher: Bencher, size: usize) {
+    bencher.bench(|| drop(BufferMut::<u8>::with_capacity(size).freeze()));
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_freeze_drop_vortex_custom(bencher: Bencher, size: usize) {
+    bencher
+        .with_inputs(|| BufferAllocatorRef::new(Global))
+        .bench_refs(|allocator| drop(allocator.with_capacity::<u8>(size).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(),
+        )
+    });
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_freeze_drop_bytes(bencher: Bencher, size: usize) {
+    bencher.bench(|| drop(BytesMut::with_capacity(size).freeze()));
+}
+
+#[divan::bench(args = SIZES)]
+fn allocate_freeze_drop_arrow(bencher: Bencher, size: usize) {
+    bencher.bench(|| {
+        let buffer: arrow_buffer::Buffer = MutableBuffer::with_capacity(size).into();
+        drop(buffer)
+    });
+}
+
+#[divan::bench(args = SIZES)]
+fn from_vec_drop_vortex(bencher: Bencher, size: usize) {
+    bencher
+        .with_inputs(|| vec![0u8; size])
+        .bench_values(|values| drop(Buffer::from(values)));
+}
+
+#[divan::bench(args = SIZES)]
+fn from_vec_drop_bytes(bencher: Bencher, size: usize) {
+    bencher
+        .with_inputs(|| vec![0u8; size])
+        .bench_values(|values| drop(bytes::Bytes::from(values)));
+}
+
+#[divan::bench(args = SIZES)]
+fn from_vec_drop_arrow(bencher: Bencher, size: usize) {
+    bencher
+        .with_inputs(|| vec![0u8; size])
+        .bench_values(|values| drop(arrow_buffer::Buffer::from_vec(values)));
+}