Xuanwo commented on code in PR #6336:
URL: https://github.com/apache/arrow-rs/pull/6336#discussion_r1738171391
##########
arrow-buffer/src/buffer/mutable.rs:
##########
@@ -483,6 +499,139 @@ 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 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"))]
+ {
+ 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
+ return;
+ }
+ 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;
+ return;
+ }
+ 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.data = NonNull::new(data).unwrap_or_else(||
handle_alloc_error(new_layout));
Review Comment:
unused code?
##########
arrow-buffer/src/buffer/mutable.rs:
##########
@@ -483,6 +499,139 @@ 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 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"))]
+ {
+ 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")]
Review Comment:
The entire `MutableBuffer` has been gated under `allocator_api`. Do we still
need it?
##########
arrow-buffer/src/buffer/mutable.rs:
##########
@@ -28,6 +31,23 @@ use crate::{
use super::Buffer;
+#[cfg(not(feature = "allocator_api"))]
+pub trait Allocator: private::Sealed {}
Review Comment:
At first glance, it's somewhat confusing to define a sealed trait here.
Maybe add some comments here?
##########
arrow-buffer/src/buffer/mutable.rs:
##########
@@ -483,6 +499,139 @@ 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"))]
+ {
Review Comment:
How about adding a `let _ = alloc` here instead of using `_alloc`?
##########
arrow-buffer/src/buffer/mutable.rs:
##########
@@ -28,6 +31,23 @@ use crate::{
use super::Buffer;
+#[cfg(not(feature = "allocator_api"))]
Review Comment:
I'm not sure if it's a good idea to introduce
[allocator-api2](https://github.com/zakarumych/allocator-api2) as a
compatibility layer?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
