waynexia commented on code in PR #6336:
URL: https://github.com/apache/arrow-rs/pull/6336#discussion_r1738741765
##########
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:
You may misread this. Only those new public APIs are limited under the
feature gate.
--
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]
