Mutablebuffer::shrink_to_fit

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);
+    }
 }