feat: add Allocator type param to MutableBuffer

.github/workflows/arrow.yml

@@ -61,8 +61,8 @@ jobs:
           submodules: true
       - name: Setup Rust toolchain
         uses: ./.github/actions/setup-builder
-      - name: Test arrow-buffer with all features
-        run: cargo test -p arrow-buffer --all-features
+      - name: Test arrow-buffer
+        run: cargo test -p arrow-buffer
       - name: Test arrow-data with all features
         run: cargo test -p arrow-data --all-features
       - name: Test arrow-schema with all features
@@ -163,8 +163,8 @@ jobs:
         uses: ./.github/actions/setup-builder
       - name: Setup Clippy
         run: rustup component add clippy
-      - name: Clippy arrow-buffer with all features
-        run: cargo clippy -p arrow-buffer --all-targets --all-features -- -D warnings
+      - name: Clippy arrow-buffer
+        run: cargo clippy -p arrow-buffer --all-targets -- -D warnings
       - name: Clippy arrow-data with all features
         run: cargo clippy -p arrow-data --all-targets --all-features -- -D warnings
       - name: Clippy arrow-schema with all features

.github/workflows/miri.sh

@@ -12,7 +12,7 @@ cargo miri setup
 cargo clean
 
 echo "Starting Arrow MIRI run..."
-cargo miri test -p arrow-buffer
+cargo miri test -p arrow-buffer --features allocator_api
 cargo miri test -p arrow-data --features ffi
 cargo miri test -p arrow-schema --features ffi
 cargo miri test -p arrow-ord

arrow-buffer/Cargo.toml

@@ -33,6 +33,9 @@ name = "arrow_buffer"
 path = "src/lib.rs"
 bench = false
 
+[features]
+allocator_api = []
+
 [dependencies]
 bytes = { version = "1.4" }
 num = { version = "0.4", default-features = false, features = ["std"] }

arrow-buffer/src/buffer/mutable.rs

@@ -15,6 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#[cfg(feature = "allocator_api")]
+use std::alloc::{Allocator, Global};
+
 use std::alloc::{handle_alloc_error, Layout};
 use std::mem;
 use std::ptr::NonNull;
@@ -28,6 +31,25 @@ use crate::{
 
 use super::Buffer;
 
+#[cfg(not(feature = "allocator_api"))]
+#[doc = "Placeholder trait for std::alloc::Allocator. To appoint a real allocator,"]
+#[doc = "please enable `allocator_api` feature with nightly toolchain"]
+pub trait Allocator: private::Sealed {}
+
+#[cfg(not(feature = "allocator_api"))]
+impl Allocator for Global {}
+
+#[cfg(not(feature = "allocator_api"))]
+#[derive(Debug)]
+pub struct Global;
+
+#[cfg(not(feature = "allocator_api"))]
+mod private {
+    pub trait Sealed {}
+
+    impl Sealed for super::Global {}
+}
+
 /// A [`MutableBuffer`] is Arrow's interface to build a [`Buffer`] out of items or slices of items.
 ///
 /// [`Buffer`]s created from [`MutableBuffer`] (via `into`) are guaranteed to have its pointer aligned
@@ -51,12 +73,20 @@ use super::Buffer;
 /// assert_eq!(buffer.as_slice(), &[0u8, 1, 0, 0, 1, 0, 0, 0])
 /// ```
 #[derive(Debug)]
-pub struct MutableBuffer {
+pub struct MutableBuffer<
+    #[cfg(feature = "allocator_api")] A: Allocator = Global,
+    #[cfg(not(feature = "allocator_api"))] A: Allocator = Global,
+> {
     // dangling iff capacity = 0
     data: NonNull<u8>,
     // invariant: len <= capacity
     len: usize,
     layout: Layout,
+    #[cfg(feature = "allocator_api")]
+    allocator: A,
+    #[cfg(not(feature = "allocator_api"))]
+    #[doc = "Placeholder for allocator API"]
+    allocator: A,
 }
 
 impl MutableBuffer {
@@ -83,14 +113,14 @@ impl MutableBuffer {
             0 => dangling_ptr(),
             _ => {
                 // Safety: Verified size != 0
-                let raw_ptr = unsafe { std::alloc::alloc(layout) };
-                NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout))
+                unsafe { Self::alloc(&Global, layout) }
             }
         };
         Self {
             data,
             len: 0,
             layout,
+            allocator: Global,
         }
     }
 
@@ -115,7 +145,12 @@ impl MutableBuffer {
                 NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout))
             }
         };
-        Self { data, len, layout }
+        Self {
+            data,
+            len,
+            layout,
+            allocator: Global,
+        }
     }
 
     /// Create a [`MutableBuffer`] from the provided [`Vec`] without copying
@@ -136,7 +171,12 @@ impl MutableBuffer {
         let data = bytes.ptr();
         mem::forget(bytes);
 
-        Ok(Self { data, len, layout })
+        Ok(Self {
+            data,
+            len,
+            layout,
+            allocator: Global,
+        })
     }
 
     /// creates a new [MutableBuffer] with capacity and length capable of holding `len` bits.
@@ -204,28 +244,6 @@ impl MutableBuffer {
         }
     }
 
-    #[cold]
-    fn reallocate(&mut self, capacity: usize) {
-        let new_layout = Layout::from_size_align(capacity, self.layout.align()).unwrap();
-        if new_layout.size() == 0 {
-            if self.layout.size() != 0 {
-                // Safety: data was allocated with layout
-                unsafe { std::alloc::dealloc(self.as_mut_ptr(), self.layout) };
-                self.layout = new_layout
-            }
-            return;
-        }
-
-        let data = match self.layout.size() {
-            // Safety: new_layout is not empty
-            0 => unsafe { std::alloc::alloc(new_layout) },
-            // Safety: verified new layout is valid and not empty
-            _ => unsafe { std::alloc::realloc(self.as_mut_ptr(), self.layout, capacity) },
-        };
-        self.data = NonNull::new(data).unwrap_or_else(|| handle_alloc_error(new_layout));
-        self.layout = new_layout;
-    }
-
     /// Truncates this buffer to `len` bytes
     ///
     /// If `len` is greater than the buffer's current length, this has no effect
@@ -483,6 +501,138 @@ impl MutableBuffer {
     }
 }
 
+#[cfg(feature = "allocator_api")]
+impl<A: Allocator> MutableBuffer<A> {
+    /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`
+    /// in the given allocator.
+    ///
+    /// See [`MutableBuffer::with_capacity_in`].
+    #[inline]
+    pub fn new_in(allocator: A, capacity: usize) -> Self {
+        Self::with_capacity_in(allocator, capacity)
+    }
+
+    /// Allocate a new [MutableBuffer] with initial capacity to be at least `capacity`.
+    /// in the given allocator
+    ///
+    /// # Panics
+    ///
+    /// If `capacity`, when rounded up to the nearest multiple of [`ALIGNMENT`], is greater
+    /// then `isize::MAX`, then this function will panic.
+    #[inline]
+    pub fn with_capacity_in(allocator: A, capacity: usize) -> Self {
+        let capacity = bit_util::round_upto_multiple_of_64(capacity);
+        let layout = Layout::from_size_align(capacity, ALIGNMENT)
+            .expect("failed to create layout for MutableBuffer");
+        let data = match layout.size() {
+            0 => dangling_ptr(),
+            _ => {
+                // Safety: Verified size != 0
+                unsafe { Self::alloc(&allocator, layout) }
+            }
+        };
+        Self {
+            data,
+            len: 0,
+            layout,
+            allocator,
+        }
+    }
+}
+
+/// `allocator_api` related internal methods
+impl<A: Allocator> MutableBuffer<A> {
+    #[inline]
+    unsafe fn alloc(alloc: &A, layout: Layout) -> NonNull<u8> {
+        #[cfg(feature = "allocator_api")]
+        {
+            alloc
+                .allocate(layout)
+                .unwrap_or_else(|_| handle_alloc_error(layout))
+                .cast()
+        }
+
+        #[cfg(not(feature = "allocator_api"))]
+        {
+            let _ = alloc;
+            let data = std::alloc::alloc(layout);
+            NonNull::new(data).unwrap_or_else(|| handle_alloc_error(layout))
+        }
+    }
+
+    #[inline]
+    unsafe fn dealloc(alloc: &A, ptr: NonNull<u8>, layout: Layout) {
+        #[cfg(feature = "allocator_api")]
+        {
+            alloc.deallocate(ptr, layout)
+        }
+
+        #[cfg(not(feature = "allocator_api"))]
+        {
+            let _ = alloc;
+            std::alloc::dealloc(ptr.as_ptr(), layout)
+        }
+    }
+
+    #[cold]
+    fn reallocate(&mut self, capacity: usize) {
+        let new_layout = Layout::from_size_align(capacity, self.layout.align()).unwrap();
+
+        // shrink to zero
+        if new_layout.size() == 0 {
+            if self.layout.size() != 0 {
+                // Safety: data was allocated with layout
+                unsafe { Self::dealloc(&self.allocator, self.data, self.layout) };
+                self.layout = new_layout
+            }
+            return;
+        }
+
+        #[cfg(feature = "allocator_api")]
+        match new_layout.size().cmp(&self.layout.size()) {
+            std::cmp::Ordering::Equal => {
+                // no action needed
+            }
+            std::cmp::Ordering::Less => {
+                // shrink to new capacity
+                let new_data = unsafe {
+                    self.allocator
+                        .shrink(self.data, self.layout, new_layout)
+                        .unwrap_or_else(|_| handle_alloc_error(new_layout))
+                        .cast()
+                };
+                self.layout = new_layout;
+                self.data = new_data;
+            }
+            std::cmp::Ordering::Greater => {
+                // grow to new capacity
+                let new_data = unsafe {
+                    self.allocator
+                        .grow(self.data, self.layout, new_layout)
+                        .unwrap_or_else(|_| handle_alloc_error(new_layout))
+                        .cast()
+                };
+                self.layout = new_layout;
+                self.data = new_data;
+            }
+        }
+
+        #[cfg(not(feature = "allocator_api"))]
+        {
+            self.data = match self.layout.size() {
+                // Safety: new_layout is not empty
+                0 => unsafe { Self::alloc(&self.allocator, new_layout) },
+                // Safety: verified new layout is valid and not empty
+                _ => unsafe {
+                    let new_data = std::alloc::realloc(self.data.as_ptr(), self.layout, capacity);
+                    NonNull::new(new_data).unwrap_or_else(|| handle_alloc_error(new_layout))
+                },
+            };
+            self.layout = new_layout;
+        }
+    }
+}
+
 #[inline]
 fn dangling_ptr() -> NonNull<u8> {
     // SAFETY: ALIGNMENT is a non-zero usize which is then casted
@@ -517,8 +667,18 @@ impl<T: ArrowNativeType> From<Vec<T>> for MutableBuffer {
         // Vec guaranteed to have a valid layout matching that of `Layout::array`
         // This is based on `RawVec::current_memory`
         let layout = unsafe { Layout::array::<T>(value.capacity()).unwrap_unchecked() };
+        let zelf = Self {
+            data,
+            len,
+            layout,
+            #[cfg(not(feature = "allocator_api"))]
+            allocator: Global,
+            #[cfg(feature = "allocator_api")]
+            allocator: *value.allocator(),
+        };
+
         mem::forget(value);
-        Self { data, len, layout }
+        zelf
     }
 }
 
@@ -688,11 +848,11 @@ impl std::ops::DerefMut for MutableBuffer {
     }
 }
 
-impl Drop for MutableBuffer {
+impl<A: Allocator> Drop for MutableBuffer<A> {
     fn drop(&mut self) {
         if self.layout.size() != 0 {
             // Safety: data was allocated with standard allocator with given layout
-            unsafe { std::alloc::dealloc(self.data.as_ptr() as _, self.layout) };
+            unsafe { Self::dealloc(&self.allocator, self.data, self.layout) };
         }
     }
 }

arrow-buffer/src/builder/null.rs

@@ -18,6 +18,7 @@
 use crate::{BooleanBufferBuilder, MutableBuffer, NullBuffer};
 
 /// Builder for creating the null bit buffer.
+///
 /// This builder only materializes the buffer when we append `false`.
 /// If you only append `true`s to the builder, what you get will be
 /// `None` when calling [`finish`](#method.finish).

arrow-buffer/src/lib.rs

@@ -19,6 +19,7 @@
 
 // used by [`buffer::mutable::dangling_ptr`]
 #![cfg_attr(miri, feature(strict_provenance))]
+#![cfg_attr(feature = "allocator_api", feature(allocator_api))]
 
 pub mod alloc;
 pub mod buffer;

arrow-buffer/src/util/bit_mask.rs

@@ -20,9 +20,11 @@
 use crate::bit_chunk_iterator::BitChunks;
 use crate::bit_util::{ceil, get_bit, set_bit};
 
-/// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]` to be equal to the
-/// bits in `data` in the range `[offset_read..offset_read+len]`
-/// returns the number of `0` bits `data[offset_read..offset_read+len]`
+/// Sets bits by range.
+///
+/// Sets all bits on `write_data` in the range `[offset_write..offset_write+len]`
+/// to be equal to the bits in `data` in the range `[offset_read..offset_read+len]`.
+/// Returns the number of `0` bits `data[offset_read..offset_read+len]`.
 pub fn set_bits(
     write_data: &mut [u8],
     data: &[u8],