[arrow-rs] branch active_release updated: Mutablebuffer::shrink_to_fit (#318) (#344)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch active_release
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/active_release by this push:
     new 305b41e  Mutablebuffer::shrink_to_fit (#318) (#344)
305b41e is described below

commit 305b41e97692b0b342846c3bff06e42f439772fa
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Tue May 25 11:30:59 2021 -0400

    Mutablebuffer::shrink_to_fit (#318) (#344)
    
    * Mutablebuffer::shrink_to_fit
    
    * add shrink_to_fit explicit test
    
    Co-authored-by: Ritchie Vink <ri...@gmail.com>
---
 arrow/src/buffer/mutable.rs | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/arrow/src/buffer/mutable.rs b/arrow/src/buffer/mutable.rs
index d7fd5b9..35b123d 100644
--- a/arrow/src/buffer/mutable.rs
+++ b/arrow/src/buffer/mutable.rs
@@ -176,6 +176,37 @@ impl MutableBuffer {
         self.len = new_len;
     }
 
+    /// Shrinks the capacity of the buffer as much as possible.
+    /// The new capacity will aligned to the nearest 64 bit alignment.
+    ///
+    /// # Example
+    /// ```
+    /// # use arrow::buffer::{Buffer, MutableBuffer};
+    /// // 2 cache lines
+    /// let mut buffer = MutableBuffer::new(128);
+    /// assert_eq!(buffer.capacity(), 128);
+    /// buffer.push(1);
+    /// buffer.push(2);
+    ///
+    /// buffer.shrink_to_fit();
+    /// assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
+    /// ```
+    pub fn shrink_to_fit(&mut self) {
+        let new_capacity = bit_util::round_upto_multiple_of_64(self.len);
+        if new_capacity < self.capacity {
+            // JUSTIFICATION
+            //  Benefit
+            //      necessity
+            //  Soundness
+            //      `self.data` is valid for `self.capacity`.
+            let ptr =
+                unsafe { alloc::reallocate(self.data, self.capacity, new_capacity) };
+
+            self.data = ptr;
+            self.capacity = new_capacity;
+        }
+    }
+
     /// Returns whether this buffer is empty or not.
     #[inline]
     pub const fn is_empty(&self) -> bool {
@@ -746,4 +777,15 @@ mod tests {
         buf2.reserve(65);
         assert!(buf != buf2);
     }
+
+    #[test]
+    fn test_mutable_shrink_to_fit() {
+        let mut buffer = MutableBuffer::new(128);
+        assert_eq!(buffer.capacity(), 128);
+        buffer.push(1);
+        buffer.push(2);
+
+        buffer.shrink_to_fit();
+        assert!(buffer.capacity() >= 64 && buffer.capacity() < 128);
+    }
 }