From 42d7e4f514d6137806dd779941b1ae4d6d0d04a0 Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Tue, 18 May 2021 15:00:52 +0200 Subject: [PATCH 1/2] Mutablebuffer::shrink_to_fit --- arrow/src/buffer/mutable.rs | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/arrow/src/buffer/mutable.rs b/arrow/src/buffer/mutable.rs index d7fd5b9d2002..eb8f5b4e82f1 100644 --- a/arrow/src/buffer/mutable.rs +++ b/arrow/src/buffer/mutable.rs @@ -176,6 +176,40 @@ 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); + /// ``` + // For performance reasons, this must be inlined so that the `if` is executed inside the caller, and not as an extra call that just + // exits. + #[inline(always)] + 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 { From 830d3a8066448a5363eb82cf3b8fd1a2571b65ee Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Tue, 18 May 2021 18:26:39 +0200 Subject: [PATCH 2/2] add shrink_to_fit explicit test --- arrow/src/buffer/mutable.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/arrow/src/buffer/mutable.rs b/arrow/src/buffer/mutable.rs index eb8f5b4e82f1..35b123d1fc44 100644 --- a/arrow/src/buffer/mutable.rs +++ b/arrow/src/buffer/mutable.rs @@ -191,9 +191,6 @@ impl MutableBuffer { /// buffer.shrink_to_fit(); /// assert!(buffer.capacity() >= 64 && buffer.capacity() < 128); /// ``` - // For performance reasons, this must be inlined so that the `if` is executed inside the caller, and not as an extra call that just - // exits. - #[inline(always)] pub fn shrink_to_fit(&mut self) { let new_capacity = bit_util::round_upto_multiple_of_64(self.len); if new_capacity < self.capacity { @@ -780,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); + } }