emilk commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2500094772
I've added a test to make sure the `shrink_to_fit` call actually returns
memory to global allocator, as is its job.
This works with the current code, but if I try to use `std::alloc::realloc`
the memory is never returned. It is as if the `realloc` call succeeds without
actually freeing up the extra memory. Here is the diff for what I tried:
```diff
diff --git a/arrow-buffer/src/buffer/immutable.rs
b/arrow-buffer/src/buffer/immutable.rs
index 40625329a..70a100734 100644
--- a/arrow-buffer/src/buffer/immutable.rs
+++ b/arrow-buffer/src/buffer/immutable.rs
@@ -173,7 +173,15 @@ impl Buffer {
///
/// If the capacity is already less than or equal to the desired
capacity, this is a no-op.
pub fn shrink_to_fit(&mut self) {
- if self.len() < self.capacity() {
+ let desired_capacity = self.len();
+ if desired_capacity < self.capacity() {
+ if let Some(bytes) = Arc::get_mut(&mut self.data) {
+ if bytes.try_realloc(desired_capacity).is_ok() {
+ return;
+ }
+ }
+
+ // Fallback:
*self = Self::from_vec(self.as_slice().to_vec())
}
}
diff --git a/arrow-buffer/src/bytes.rs b/arrow-buffer/src/bytes.rs
index ba61342d8..e616ee7c0 100644
--- a/arrow-buffer/src/bytes.rs
+++ b/arrow-buffer/src/bytes.rs
@@ -96,6 +96,27 @@ impl Bytes {
}
}
+ /// Try to reallocate the underlying memory region to a new size
(smaller or larger).
+ ///
+ /// Returns `Err(())` if the reallocation failed.
+ /// Only works for bytes allocated with the standard allocator.
+ pub fn try_realloc(&mut self, new_len: usize) -> Result<(), ()> {
+ if let Deallocation::Standard(layout) = self.deallocation {
+ if let Ok(new_layout) =
std::alloc::Layout::from_size_align(new_len, layout.align()) {
+ let new_ptr =
+ unsafe { std::alloc::realloc(self.ptr.as_mut(),
new_layout, new_len) };
+ if let Some(ptr) = NonNull::new(new_ptr) {
+ self.ptr = ptr;
+ self.len = new_len;
+ self.deallocation = Deallocation::Standard(new_layout);
+ return Ok(());
+ }
+ }
+ }
+
+ Err(())
+ }
+
#[inline]
pub(crate) fn deallocation(&self) -> &Deallocation {
&self.deallocation
```
--
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]
