Add Array::shrink_to_fit(&mut self) to 53.4.0 (#6790) (#6817) (#6962)

emilk · tustvold · web-flow · commit 58f086fc2859 · 2025-01-13T12:35:30.000-05:00
* Add `Array::shrink_to_fit(&mut self)` (#6790) * Add `Array::shrink_to_fit` * Test that shrink_to_fit actually frees memory * Make sure the buffer isn't shared in the test of shrink_to_fit * Remove `#[inline]` * Use `realloc` to reallocate the bytes * Clean up test * Improve docstring for `Array::shrink_to_fit` Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com> * `Buffer::shrink_to_fit`: ignore shared buffers * Improve comment in `ArrayRef::shrink_to_fit` * Document why `try_realloc` is safe, and actually make it safe :) * Improve testing of shrink_to_fit * Fix a few corner cases, and improve test * Add license header to new test file --------- Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com> * Support shrink to empty (#6817) * Support shrink to empty * Docs * Format --------- Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
diff --git a/arrow-array/src/array/boolean_array.rs b/arrow-array/src/array/boolean_array.rs
@@ -308,6 +308,13 @@ impl Array for BooleanArray {
         self.values.is_empty()
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+    }
+
     fn offset(&self) -> usize {
         self.values.offset()
     }
diff --git a/arrow-array/src/array/byte_array.rs b/arrow-array/src/array/byte_array.rs
@@ -453,6 +453,14 @@ impl<T: ByteArrayType> Array for GenericByteArray<T> {
         self.value_offsets.len() <= 1
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.value_offsets.shrink_to_fit();
+        self.value_data.shrink_to_fit();
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/byte_view_array.rs b/arrow-array/src/array/byte_view_array.rs
@@ -430,31 +430,31 @@ impl<T: ByteViewType + ?Sized> GenericByteViewArray<T> {
     ///
     /// Before GC:
     /// ```text
-    ///                                        ┌──────┐                 
-    ///                                        │......│                 
-    ///                                        │......│                 
-    /// ┌────────────────────┐       ┌ ─ ─ ─ ▶ │Data1 │   Large buffer  
+    ///                                        ┌──────┐
+    ///                                        │......│
+    ///                                        │......│
+    /// ┌────────────────────┐       ┌ ─ ─ ─ ▶ │Data1 │   Large buffer
     /// │       View 1       │─ ─ ─ ─          │......│  with data that
     /// ├────────────────────┤                 │......│ is not referred
     /// │       View 2       │─ ─ ─ ─ ─ ─ ─ ─▶ │Data2 │ to by View 1 or
-    /// └────────────────────┘                 │......│      View 2     
-    ///                                        │......│                 
-    ///    2 views, refer to                   │......│                 
-    ///   small portions of a                  └──────┘                 
-    ///      large buffer                                               
+    /// └────────────────────┘                 │......│      View 2
+    ///                                        │......│
+    ///    2 views, refer to                   │......│
+    ///   small portions of a                  └──────┘
+    ///      large buffer
     /// ```
-    ///                                                                
+    ///
     /// After GC:
     ///
     /// ```text
     /// ┌────────────────────┐                 ┌─────┐    After gc, only
-    /// │       View 1       │─ ─ ─ ─ ─ ─ ─ ─▶ │Data1│     data that is  
-    /// ├────────────────────┤       ┌ ─ ─ ─ ▶ │Data2│    pointed to by  
-    /// │       View 2       │─ ─ ─ ─          └─────┘     the views is  
-    /// └────────────────────┘                                 left      
-    ///                                                                  
-    ///                                                                  
-    ///         2 views                                                  
+    /// │       View 1       │─ ─ ─ ─ ─ ─ ─ ─▶ │Data1│     data that is
+    /// ├────────────────────┤       ┌ ─ ─ ─ ▶ │Data2│    pointed to by
+    /// │       View 2       │─ ─ ─ ─          └─────┘     the views is
+    /// └────────────────────┘                                 left
+    ///
+    ///
+    ///         2 views
     /// ```
     /// This method will compact the data buffers by recreating the view array and only include the data
     /// that is pointed to by the views.
@@ -575,6 +575,15 @@ impl<T: ByteViewType + ?Sized> Array for GenericByteViewArray<T> {
         self.views.is_empty()
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.views.shrink_to_fit();
+        self.buffers.iter_mut().for_each(|b| b.shrink_to_fit());
+        self.buffers.shrink_to_fit();
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/dictionary_array.rs b/arrow-array/src/array/dictionary_array.rs
@@ -720,6 +720,11 @@ impl<T: ArrowDictionaryKeyType> Array for DictionaryArray<T> {
         self.keys.is_empty()
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.keys.shrink_to_fit();
+        self.values.shrink_to_fit();
+    }
+
     fn offset(&self) -> usize {
         self.keys.offset()
     }
diff --git a/arrow-array/src/array/fixed_size_binary_array.rs b/arrow-array/src/array/fixed_size_binary_array.rs
@@ -602,6 +602,13 @@ impl Array for FixedSizeBinaryArray {
         self.len == 0
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.value_data.shrink_to_fit();
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/fixed_size_list_array.rs b/arrow-array/src/array/fixed_size_list_array.rs
@@ -401,6 +401,13 @@ impl Array for FixedSizeListArray {
         self.len == 0
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs
@@ -485,6 +485,14 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericListArray<OffsetSize> {
         self.value_offsets.len() <= 1
     }
 
+    fn shrink_to_fit(&mut self) {
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+        self.values.shrink_to_fit();
+        self.value_offsets.shrink_to_fit();
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/list_view_array.rs b/arrow-array/src/array/list_view_array.rs
@@ -326,6 +326,15 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericListViewArray<OffsetSize> {
         self.value_sizes.is_empty()
     }
 
+    fn shrink_to_fit(&mut self) {
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+        self.values.shrink_to_fit();
+        self.value_offsets.shrink_to_fit();
+        self.value_sizes.shrink_to_fit();
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/map_array.rs b/arrow-array/src/array/map_array.rs
@@ -372,6 +372,14 @@ impl Array for MapArray {
         self.value_offsets.len() <= 1
     }
 
+    fn shrink_to_fit(&mut self) {
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+        self.entries.shrink_to_fit();
+        self.value_offsets.shrink_to_fit();
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/mod.rs b/arrow-array/src/array/mod.rs
@@ -167,6 +167,12 @@ pub trait Array: std::fmt::Debug + Send + Sync {
     /// ```
     fn is_empty(&self) -> bool;
 
+    /// Shrinks the capacity of any exclusively owned buffer as much as possible
+    ///
+    /// Shared or externally allocated buffers will be ignored, and
+    /// any buffer offsets will be preserved.
+    fn shrink_to_fit(&mut self) {}
+
     /// Returns the offset into the underlying data used by this array(-slice).
     /// Note that the underlying data can be shared by many arrays.
     /// This defaults to `0`.
@@ -366,6 +372,15 @@ impl Array for ArrayRef {
         self.as_ref().is_empty()
     }
 
+    /// For shared buffers, this is a no-op.
+    fn shrink_to_fit(&mut self) {
+        if let Some(slf) = Arc::get_mut(self) {
+            slf.shrink_to_fit();
+        } else {
+            // We ignore shared buffers.
+        }
+    }
+
     fn offset(&self) -> usize {
         self.as_ref().offset()
     }
diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
@@ -1152,6 +1152,13 @@ impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {
         self.values.is_empty()
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.values.shrink_to_fit();
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/run_array.rs b/arrow-array/src/array/run_array.rs
@@ -330,6 +330,11 @@ impl<T: RunEndIndexType> Array for RunArray<T> {
         self.run_ends.is_empty()
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.run_ends.shrink_to_fit();
+        self.values.shrink_to_fit();
+    }
+
     fn offset(&self) -> usize {
         self.run_ends.offset()
     }
diff --git a/arrow-array/src/array/struct_array.rs b/arrow-array/src/array/struct_array.rs
@@ -370,6 +370,13 @@ impl Array for StructArray {
         self.len == 0
     }
 
+    fn shrink_to_fit(&mut self) {
+        if let Some(nulls) = &mut self.nulls {
+            nulls.shrink_to_fit();
+        }
+        self.fields.iter_mut().for_each(|n| n.shrink_to_fit());
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-array/src/array/union_array.rs b/arrow-array/src/array/union_array.rs
@@ -744,6 +744,17 @@ impl Array for UnionArray {
         self.type_ids.is_empty()
     }
 
+    fn shrink_to_fit(&mut self) {
+        self.type_ids.shrink_to_fit();
+        if let Some(offsets) = &mut self.offsets {
+            offsets.shrink_to_fit();
+        }
+        for array in self.fields.iter_mut().flatten() {
+            array.shrink_to_fit();
+        }
+        self.fields.shrink_to_fit();
+    }
+
     fn offset(&self) -> usize {
         0
     }
diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs
@@ -125,6 +125,12 @@ impl BooleanBuffer {
         self.len == 0
     }
 
+    /// Free up unused memory.
+    pub fn shrink_to_fit(&mut self) {
+        // TODO(emilk): we could shrink even more in the case where we are a small sub-slice of the full buffer
+        self.buffer.shrink_to_fit();
+    }
+
     /// Returns the boolean value at index `i`.
     ///
     /// # Panics
diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs
@@ -167,6 +167,41 @@ impl Buffer {
         self.data.capacity()
     }
 
+    /// Tried to shrink the capacity of the buffer as much as possible, freeing unused memory.
+    ///
+    /// If the buffer is shared, this is a no-op.
+    ///
+    /// If the memory was allocated with a custom allocator, this is a no-op.
+    ///
+    /// If the capacity is already less than or equal to the desired capacity, this is a no-op.
+    ///
+    /// The memory region will be reallocated using `std::alloc::realloc`.
+    pub fn shrink_to_fit(&mut self) {
+        let offset = self.ptr_offset();
+        let is_empty = self.is_empty();
+        let desired_capacity = if is_empty {
+            0
+        } else {
+            // For realloc to work, we cannot free the elements before the offset
+            offset + 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() {
+                    // Realloc complete - update our pointer into `bytes`:
+                    self.ptr = if is_empty {
+                        bytes.as_ptr()
+                    } else {
+                        // SAFETY: we kept all elements leading up to the offset
+                        unsafe { bytes.as_ptr().add(offset) }
+                    }
+                } else {
+                    // Failure to reallocate is fine; we just failed to free up memory.
+                }
+            }
+        }
+    }
+
     /// Returns whether the buffer is empty.
     #[inline]
     pub fn is_empty(&self) -> bool {
@@ -562,6 +597,34 @@ mod tests {
         assert_eq!(buf2.slice_with_length(2, 1).as_slice(), &[10]);
     }
 
+    #[test]
+    fn test_shrink_to_fit() {
+        let original = Buffer::from(&[0, 1, 2, 3, 4, 5, 6, 7]);
+        assert_eq!(original.as_slice(), &[0, 1, 2, 3, 4, 5, 6, 7]);
+        assert_eq!(original.capacity(), 64);
+
+        let slice = original.slice_with_length(2, 3);
+        drop(original); // Make sure the buffer isn't shared (or shrink_to_fit won't work)
+        assert_eq!(slice.as_slice(), &[2, 3, 4]);
+        assert_eq!(slice.capacity(), 64);
+
+        let mut shrunk = slice;
+        shrunk.shrink_to_fit();
+        assert_eq!(shrunk.as_slice(), &[2, 3, 4]);
+        assert_eq!(shrunk.capacity(), 5); // shrink_to_fit is allowed to keep the elements before the offset
+
+        // Test that we can handle empty slices:
+        let empty_slice = shrunk.slice_with_length(1, 0);
+        drop(shrunk); // Make sure the buffer isn't shared (or shrink_to_fit won't work)
+        assert_eq!(empty_slice.as_slice(), &[]);
+        assert_eq!(empty_slice.capacity(), 5);
+
+        let mut shrunk_empty = empty_slice;
+        shrunk_empty.shrink_to_fit();
+        assert_eq!(shrunk_empty.as_slice(), &[]);
+        assert_eq!(shrunk_empty.capacity(), 0);
+    }
+
     #[test]
     #[should_panic(expected = "the offset of the new Buffer cannot exceed the existing length")]
     fn test_slice_offset_out_of_bound() {
diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs
@@ -483,10 +483,13 @@ impl MutableBuffer {
     }
 }
 
+/// Creates a non-null pointer with alignment of [`ALIGNMENT`]
+///
+/// This is similar to [`NonNull::dangling`]
 #[inline]
-fn dangling_ptr() -> NonNull<u8> {
-    // SAFETY: ALIGNMENT is a non-zero usize which is then casted
-    // to a *mut T. Therefore, `ptr` is not null and the conditions for
+pub(crate) fn dangling_ptr() -> NonNull<u8> {
+    // SAFETY: ALIGNMENT is a non-zero usize which is then cast
+    // to a *mut u8. Therefore, `ptr` is not null and the conditions for
     // calling new_unchecked() are respected.
     #[cfg(miri)]
     {
diff --git a/arrow-buffer/src/buffer/null.rs b/arrow-buffer/src/buffer/null.rs
@@ -130,6 +130,11 @@ impl NullBuffer {
         self.buffer.is_empty()
     }
 
+    /// Free up unused memory.
+    pub fn shrink_to_fit(&mut self) {
+        self.buffer.shrink_to_fit();
+    }
+
     /// Returns the null count for this [`NullBuffer`]
     #[inline]
     pub fn null_count(&self) -> usize {
diff --git a/arrow-buffer/src/buffer/offset.rs b/arrow-buffer/src/buffer/offset.rs
@@ -133,6 +133,11 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
         Self(out.into())
     }
 
+    /// Free up unused memory.
+    pub fn shrink_to_fit(&mut self) {
+        self.0.shrink_to_fit();
+    }
+
     /// Returns the inner [`ScalarBuffer`]
     pub fn inner(&self) -> &ScalarBuffer<O> {
         &self.0
diff --git a/arrow-buffer/src/buffer/run.rs b/arrow-buffer/src/buffer/run.rs
diff --git a/arrow-buffer/src/buffer/scalar.rs b/arrow-buffer/src/buffer/scalar.rs
diff --git a/arrow-buffer/src/bytes.rs b/arrow-buffer/src/bytes.rs
diff --git a/arrow/tests/shrink_to_fit.rs b/arrow/tests/shrink_to_fit.rs

Original file line number	Diff line number	Diff line change
`@@ -308,6 +308,13 @@ impl Array for BooleanArray {`
`308`	`308`	`self.values.is_empty()`
`309`	`309`	`}`
`310`	`310`
	`311`	`+ fn shrink_to_fit(&mut self) {`
	`312`	`+ self.values.shrink_to_fit();`
	`313`	`+ if let Some(nulls) = &mut self.nulls {`
	`314`	`+ nulls.shrink_to_fit();`
	`315`	`+ }`
	`316`	`+ }`
	`317`	`+`
`311`	`318`	`fn offset(&self) -> usize {`
`312`	`319`	`self.values.offset()`
`313`	`320`	`}`
Original file line number	Diff line number	Diff line change
`@@ -453,6 +453,14 @@ impl<T: ByteArrayType> Array for GenericByteArray<T> {`
`453`	`453`	`self.value_offsets.len() <= 1`
`454`	`454`	`}`
`455`	`455`
	`456`	`+ fn shrink_to_fit(&mut self) {`
	`457`	`+ self.value_offsets.shrink_to_fit();`
	`458`	`+ self.value_data.shrink_to_fit();`
	`459`	`+ if let Some(nulls) = &mut self.nulls {`
	`460`	`+ nulls.shrink_to_fit();`
	`461`	`+ }`
	`462`	`+ }`
	`463`	`+`
`456`	`464`	`fn offset(&self) -> usize {`
`457`	`465`	`0`
`458`	`466`	`}`
Original file line number	Diff line number	Diff line change
`@@ -720,6 +720,11 @@ impl<T: ArrowDictionaryKeyType> Array for DictionaryArray<T> {`
`720`	`720`	`self.keys.is_empty()`
`721`	`721`	`}`
`722`	`722`
	`723`	`+ fn shrink_to_fit(&mut self) {`
	`724`	`+ self.keys.shrink_to_fit();`
	`725`	`+ self.values.shrink_to_fit();`
	`726`	`+ }`
	`727`	`+`
`723`	`728`	`fn offset(&self) -> usize {`
`724`	`729`	`self.keys.offset()`
`725`	`730`	`}`
Original file line number	Diff line number	Diff line change
`@@ -602,6 +602,13 @@ impl Array for FixedSizeBinaryArray {`
`602`	`602`	`self.len == 0`
`603`	`603`	`}`
`604`	`604`
	`605`	`+ fn shrink_to_fit(&mut self) {`
	`606`	`+ self.value_data.shrink_to_fit();`
	`607`	`+ if let Some(nulls) = &mut self.nulls {`
	`608`	`+ nulls.shrink_to_fit();`
	`609`	`+ }`
	`610`	`+ }`
	`611`	`+`
`605`	`612`	`fn offset(&self) -> usize {`
`606`	`613`	`0`
`607`	`614`	`}`
Original file line number	Diff line number	Diff line change
`@@ -401,6 +401,13 @@ impl Array for FixedSizeListArray {`
`401`	`401`	`self.len == 0`
`402`	`402`	`}`
`403`	`403`
	`404`	`+ fn shrink_to_fit(&mut self) {`
	`405`	`+ self.values.shrink_to_fit();`
	`406`	`+ if let Some(nulls) = &mut self.nulls {`
	`407`	`+ nulls.shrink_to_fit();`
	`408`	`+ }`
	`409`	`+ }`
	`410`	`+`
`404`	`411`	`fn offset(&self) -> usize {`
`405`	`412`	`0`
`406`	`413`	`}`
Original file line number	Diff line number	Diff line change
`@@ -485,6 +485,14 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericListArray<OffsetSize> {`
`485`	`485`	`self.value_offsets.len() <= 1`
`486`	`486`	`}`
`487`	`487`
	`488`	`+ fn shrink_to_fit(&mut self) {`
	`489`	`+ if let Some(nulls) = &mut self.nulls {`
	`490`	`+ nulls.shrink_to_fit();`
	`491`	`+ }`
	`492`	`+ self.values.shrink_to_fit();`
	`493`	`+ self.value_offsets.shrink_to_fit();`
	`494`	`+ }`
	`495`	`+`
`488`	`496`	`fn offset(&self) -> usize {`
`489`	`497`	`0`
`490`	`498`	`}`
Original file line number	Diff line number	Diff line change
`@@ -326,6 +326,15 @@ impl<OffsetSize: OffsetSizeTrait> Array for GenericListViewArray<OffsetSize> {`
`326`	`326`	`self.value_sizes.is_empty()`
`327`	`327`	`}`
`328`	`328`
	`329`	`+ fn shrink_to_fit(&mut self) {`
	`330`	`+ if let Some(nulls) = &mut self.nulls {`
	`331`	`+ nulls.shrink_to_fit();`
	`332`	`+ }`
	`333`	`+ self.values.shrink_to_fit();`
	`334`	`+ self.value_offsets.shrink_to_fit();`
	`335`	`+ self.value_sizes.shrink_to_fit();`
	`336`	`+ }`
	`337`	`+`
`329`	`338`	`fn offset(&self) -> usize {`
`330`	`339`	`0`
`331`	`340`	`}`
Original file line number	Diff line number	Diff line change
`@@ -372,6 +372,14 @@ impl Array for MapArray {`
`372`	`372`	`self.value_offsets.len() <= 1`
`373`	`373`	`}`
`374`	`374`
	`375`	`+ fn shrink_to_fit(&mut self) {`
	`376`	`+ if let Some(nulls) = &mut self.nulls {`
	`377`	`+ nulls.shrink_to_fit();`
	`378`	`+ }`
	`379`	`+ self.entries.shrink_to_fit();`
	`380`	`+ self.value_offsets.shrink_to_fit();`
	`381`	`+ }`
	`382`	`+`
`375`	`383`	`fn offset(&self) -> usize {`
`376`	`384`	`0`
`377`	`385`	`}`