You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "emilk (via GitHub)" <gi...@apache.org> on 2024/11/25 15:59:10 UTC
[PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
emilk opened a new pull request, #6790:
URL: https://github.com/apache/arrow-rs/pull/6790
# Which issue does this PR close?
* Closes #6360
# Rationale for this change
Concatenating many arrow buffers incrementally can lead to situations where the buffers are using much more memory than they need (their capacity is larger than their lengths).
Example:
```rs
use arrow::{
array::{Array, ArrayRef, ListArray, PrimitiveArray},
buffer::OffsetBuffer,
datatypes::{Field, UInt8Type},
};
fn main() {
let array0: PrimitiveArray<UInt8Type> = (0..200 * 300)
.map(|v| (v % 255) as u8)
.collect::<Vec<_>>()
.into();
let array0: ArrayRef = Arc::new(array0);
let (global, local) = memory_use(|| {
let concatenated = concatenate(array0.clone());
dbg!(concatenated.data_type());
eprintln!("expected: ~{}", how_many_bytes(concatenated.clone()));
concatenated
});
eprintln!("global: {global} bytes");
eprintln!("local: {local} bytes");
eprintln!("---");
let array1 = ListArray::new(
Field::new_list_field(array0.data_type().clone(), false).into(),
OffsetBuffer::from_lengths(std::iter::once(array0.len())),
array0.clone(),
None,
);
let array1: ArrayRef = Arc::new(array1);
let (global, local) = memory_use(|| {
let concatenated = concatenate(array1.clone()).shrink_to_fit();
dbg!(concatenated.data_type());
eprintln!("expected: ~{}", how_many_bytes(concatenated.clone()));
concatenated
});
eprintln!("global: {global} bytes");
eprintln!("local: {local} bytes");
}
fn concatenate(array: ArrayRef) -> ArrayRef {
let mut concatenated = array.clone();
for _ in 0..1000 {
concatenated = arrow::compute::kernels::concat::concat(&[&*concatenated, &*array]).unwrap();
}
concatenated
}
fn how_many_bytes(array: ArrayRef) -> u64 {
let mut array = array;
loop {
match array.data_type() {
arrow::datatypes::DataType::UInt8 => break,
arrow::datatypes::DataType::List(_) => {
let list = array.as_any().downcast_ref::<ListArray>().unwrap();
array = list.values().clone();
}
_ => unreachable!(),
}
}
array.len() as _
}
// --- Memory tracking ---
use std::sync::{
atomic::{AtomicUsize, Ordering::Relaxed},
Arc,
};
static LIVE_BYTES_GLOBAL: AtomicUsize = AtomicUsize::new(0);
thread_local! {
static LIVE_BYTES_IN_THREAD: AtomicUsize = const { AtomicUsize::new(0) } ;
}
pub struct TrackingAllocator {
allocator: std::alloc::System,
}
#[global_allocator]
pub static GLOBAL_ALLOCATOR: TrackingAllocator = TrackingAllocator {
allocator: std::alloc::System,
};
#[allow(unsafe_code)]
// SAFETY:
// We just do book-keeping and then let another allocator do all the actual work.
unsafe impl std::alloc::GlobalAlloc for TrackingAllocator {
#[allow(clippy::let_and_return)]
unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
LIVE_BYTES_IN_THREAD.with(|bytes| bytes.fetch_add(layout.size(), Relaxed));
LIVE_BYTES_GLOBAL.fetch_add(layout.size(), Relaxed);
// SAFETY:
// Just deferring
unsafe { self.allocator.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
LIVE_BYTES_IN_THREAD.with(|bytes| bytes.fetch_sub(layout.size(), Relaxed));
LIVE_BYTES_GLOBAL.fetch_sub(layout.size(), Relaxed);
// SAFETY:
// Just deferring
unsafe { self.allocator.dealloc(ptr, layout) };
}
}
fn live_bytes_local() -> usize {
LIVE_BYTES_IN_THREAD.with(|bytes| bytes.load(Relaxed))
}
fn live_bytes_global() -> usize {
LIVE_BYTES_GLOBAL.load(Relaxed)
}
/// Returns `(num_bytes_allocated, num_bytes_allocated_by_this_thread)`.
fn memory_use<R>(run: impl Fn() -> R) -> (usize, usize) {
let used_bytes_start_local = live_bytes_local();
let used_bytes_start_global = live_bytes_global();
let ret = run();
let bytes_used_local = live_bytes_local() - used_bytes_start_local;
let bytes_used_global = live_bytes_global() - used_bytes_start_global;
drop(ret);
(bytes_used_global, bytes_used_local)
}
```
If you run this you will see 12 MB is used for 6 MB of data.
Adding a call to the new `.shrink_to_fit();` on `concatenated` removes the memory overhead.
# What changes are included in this PR?
This PR adds `shrink_to_fit(&mu self)` to `Array` and all buffers.
It is best-effort.
# Are there any user-facing changes?
`trait Array` now has a `fn shrink_to_fit(&mut self) { }`.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2500367864
Good idea with `realloc` - thanks for the swift feedback. I think this PR is ready for a final review now!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1863801829
##########
arrow-buffer/src/buffer/immutable.rs:
##########
@@ -576,30 +591,30 @@ mod tests {
#[test]
fn test_shrink_to_fit() {
- let original = Buffer::from(&[1, 2, 3, 4, 5, 6, 7]);
- assert_eq!(original.len(), 7);
+ 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(3);
+ 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.len(), 4);
+ assert_eq!(slice.as_slice(), &[2, 3, 4]);
assert_eq!(slice.capacity(), 64);
let mut shrunk = slice;
shrunk.shrink_to_fit();
- assert_eq!(shrunk.len(), 4);
- assert_eq!(shrunk.capacity(), 4);
+ 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(4);
+ 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.len(), 0);
- assert_eq!(empty_slice.capacity(), 4);
+ 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.len(), 0);
- assert_eq!(shrunk_empty.capacity(), 1); // `Buffer` and `Bytes` doesn't support 0-capacity, so we shrink to 1
+ assert_eq!(shrunk_empty.as_slice(), &[]);
+ assert_eq!(shrunk_empty.capacity(), 1); // NOTE: `Buffer` and `Bytes` doesn't support 0-capacity
Review Comment:
Added in https://github.com/apache/arrow-rs/pull/6817
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2501100758
The release schedule can be found in the repository readme, the next release will be a major release in early December. Unfortunately we've already started integrating breaking changes and so a patch release is unlikely save for a major security vulnerability.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1858448486
##########
arrow-buffer/src/bytes.rs:
##########
@@ -96,6 +96,30 @@ impl Bytes {
}
}
+ /// Try to reallocate the underlying memory region to a new size (smaller or larger).
+ ///
+ /// Only works for bytes allocated with the standard allocator.
+ /// Returns `Err` if the memory was allocated with a custom allocator,
+ /// or the call to `realloc` failed, for whatever reason.
+ /// In case of `Err`, the [`Bytes`] will remain as it was (i.e. have the old size).
+ pub fn try_realloc(&mut self, new_len: usize) -> Result<(), ()> {
Review Comment:
I've documented why it is safe now, and while doing so found a bug where `new_len = 0` would produce UB. That is now fixed :)
I also added a test for the zero-sized case.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1856881778
##########
arrow-array/src/array/mod.rs:
##########
@@ -365,6 +368,16 @@ impl Array for ArrayRef {
self.as_ref().is_empty()
}
+ fn shrink_to_fit(&mut self) {
+ if let Some(slf) = Arc::get_mut(self) {
+ slf.shrink_to_fit();
+ } else {
+ // TODO(emilk): clone the contents and shrink that.
+ // This can be accomplished if we add `trait Array { fn clone(&self) -> Box<Array>>; }`.
+ // Or we clone using `let clone = self.slice(0, self.len());` and hope that the returned `ArrayRef` is not shared.
Review Comment:
I think any of the following three approaches makes sense here:
## Do nothing
The current approach. If the `ArrayRef` is shared, then we cannot free up memory.
In many ways this makes sense, and I'd be fine with it. We can also start here and improve it later.
## Clone using `slice`
```rs
let mut clone = self.slice(0, self.len());
if let Some(slf) = Arc::get_mut(&mut clone) {
clone.shrink_to_fit();
} else {
// This would be highly unlikely
}
clone
```
## Clone using a new trait method
Add
``` rs
trait Array {
fn clone(&self) -> Box<dyn Array>;
}
(and implement it everywhere)
```
and then call
```rs
let mut clone = Array::clone(self);
clone.shrink_to_fit();
Arc::new(clone)
```
##########
arrow-array/src/array/mod.rs:
##########
@@ -365,6 +368,16 @@ impl Array for ArrayRef {
self.as_ref().is_empty()
}
+ fn shrink_to_fit(&mut self) {
+ if let Some(slf) = Arc::get_mut(self) {
+ slf.shrink_to_fit();
+ } else {
+ // TODO(emilk): clone the contents and shrink that.
+ // This can be accomplished if we add `trait Array { fn clone(&self) -> Box<Array>>; }`.
+ // Or we clone using `let clone = self.slice(0, self.len());` and hope that the returned `ArrayRef` is not shared.
Review Comment:
I think any of the following three approaches makes sense here:
## Do nothing
The current approach. If the `ArrayRef` is shared, then we cannot free up memory.
In many ways this makes sense, and I'd be fine with it. We can also start here and improve it later.
## Clone using `slice`
```rs
let mut clone = self.slice(0, self.len());
if let Some(slf) = Arc::get_mut(&mut clone) {
clone.shrink_to_fit();
} else {
// This would be highly unlikely, and fine
}
clone
```
## Clone using a new trait method
Add
``` rs
trait Array {
fn clone(&self) -> Box<dyn Array>;
}
(and implement it everywhere)
```
and then call
```rs
let mut clone = Array::clone(self);
clone.shrink_to_fit();
Arc::new(clone)
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1857255351
##########
arrow-buffer/src/buffer/immutable.rs:
##########
@@ -167,6 +167,17 @@ impl Buffer {
self.data.capacity()
}
+ /// Shrinks the capacity of the buffer as much as possible, freeing unused memory.
+ ///
+ /// The capacity of the returned buffer will be the same as [`Self::len`].
+ ///
+ /// 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() {
+ *self = Self::from_vec(self.as_slice().to_vec())
+ }
Review Comment:
This will always create a fresh allocation and copy across.
I think this should be a no-op if the buffer is shared or externally allocated, and otherwise call through to `realloc`. In most cases this will avoid any data copying at all, as it will just be adjusting the allocator's metadata.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2508171905
Looks good to me
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1857250921
##########
arrow-array/src/array/mod.rs:
##########
@@ -365,6 +368,16 @@ impl Array for ArrayRef {
self.as_ref().is_empty()
}
+ fn shrink_to_fit(&mut self) {
+ if let Some(slf) = Arc::get_mut(self) {
+ slf.shrink_to_fit();
+ } else {
+ // TODO(emilk): clone the contents and shrink that.
+ // This can be accomplished if we add `trait Array { fn clone(&self) -> Box<Array>>; }`.
+ // Or we clone using `let clone = self.slice(0, self.len());` and hope that the returned `ArrayRef` is not shared.
Review Comment:
Doing nothing seems fine to me
##########
arrow-buffer/src/buffer/immutable.rs:
##########
@@ -167,6 +167,17 @@ impl Buffer {
self.data.capacity()
}
+ /// Shrinks the capacity of the buffer as much as possible, freeing unused memory.
+ ///
+ /// The capacity of the returned buffer will be the same as [`Self::len`].
+ ///
+ /// 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() {
+ *self = Self::from_vec(self.as_slice().to_vec())
+ }
Review Comment:
This will always create a fresh allocation and copy across.
I think this should be a no-op if the buffer is shared or externally allocated, and otherwise call through to `realloc`. In most cases this will avoid any data copying at all.
##########
arrow-buffer/src/buffer/offset.rs:
##########
@@ -133,6 +133,12 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
Self(out.into())
}
+ /// Free up unused memory.
+ #[inline]
Review Comment:
Given these methods are performing memory manipulations, I think it unlikely `inline` will yield any performance difference and may hurt compile times.
##########
arrow-array/src/array/mod.rs:
##########
@@ -167,6 +167,9 @@ pub trait Array: std::fmt::Debug + Send + Sync {
/// ```
fn is_empty(&self) -> bool;
+ /// Frees up unused memory.
Review Comment:
```suggestion
/// 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.
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2500321672
> How are you measuring this? Many allocators, glibc especially, hang onto memory that isn't in use to avoid thrashing the TLB. They will reuse it, they just won't give it back to the system.
I install an allocator and note down all calls to it. But I realized I had a bug in my call to `realloc` - I need to give it the layout of the _previous_ allocation, but I gave it the layout to the new allocation. Fixing that bug makes the realloc path work π
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1858307737
##########
arrow-buffer/src/bytes.rs:
##########
@@ -96,6 +96,30 @@ impl Bytes {
}
}
+ /// Try to reallocate the underlying memory region to a new size (smaller or larger).
+ ///
+ /// Only works for bytes allocated with the standard allocator.
+ /// Returns `Err` if the memory was allocated with a custom allocator,
+ /// or the call to `realloc` failed, for whatever reason.
+ /// In case of `Err`, the [`Bytes`] will remain as it was (i.e. have the old size).
+ pub fn try_realloc(&mut self, new_len: usize) -> Result<(), ()> {
Review Comment:
I don't think this needs to be unsafe, however, Bytes is crate private so worst case it won't be a breaking change to fix
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1856881778
##########
arrow-array/src/array/mod.rs:
##########
@@ -365,6 +368,16 @@ impl Array for ArrayRef {
self.as_ref().is_empty()
}
+ fn shrink_to_fit(&mut self) {
+ if let Some(slf) = Arc::get_mut(self) {
+ slf.shrink_to_fit();
+ } else {
+ // TODO(emilk): clone the contents and shrink that.
+ // This can be accomplished if we add `trait Array { fn clone(&self) -> Box<Array>>; }`.
+ // Or we clone using `let clone = self.slice(0, self.len());` and hope that the returned `ArrayRef` is not shared.
Review Comment:
I think any of the following three approaches makes sense here:
## Do nothing
The current approach. If the `ArrayRef` is shared, then we cannot free up memory.
In many ways this makes sense, and I'd be fine with it. We can also start here and improve it later.
## Clone using `slice`
```rs
let mut clone = self.slice(0, self.len());
if let Some(slf) = Arc::get_mut(&mut clone) {
clone.shrink_to_fit();
} else {
// This would be highly unlikely, and fine
}
clone
```
## Clone using a new trait method
Add
``` rs
trait Array {
fn clone(&self) -> Box<dyn Array>;
}
```
(and implement it everywhere)
and then in `ArrayRef::shrink_to_fit`:s `else`:
```rs
let mut clone = Array::clone(self);
clone.shrink_to_fit();
Arc::new(clone)
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1858295068
##########
arrow/tests/shrink_to_fit.rs:
##########
@@ -0,0 +1,142 @@
+use arrow::{
+ array::{Array, ArrayRef, ListArray, PrimitiveArray},
+ buffer::OffsetBuffer,
+ datatypes::{Field, UInt8Type},
+};
+
+/// Test that `shrink_to_fit` frees memory after concatenating a large number of arrays.
+#[test]
+fn test_shrink_to_fit_after_concat() {
+ let array_len = 6_000;
+ let num_concats = 100;
+
+ let primitive_array: PrimitiveArray<UInt8Type> = (0..array_len)
+ .map(|v| (v % 255) as u8)
+ .collect::<Vec<_>>()
+ .into();
+ let primitive_array: ArrayRef = Arc::new(primitive_array);
+
+ let list_array: ArrayRef = Arc::new(ListArray::new(
+ Field::new_list_field(primitive_array.data_type().clone(), false).into(),
+ OffsetBuffer::from_lengths([primitive_array.len()]),
+ primitive_array.clone(),
+ None,
+ ));
+
+ // Num bytes allocated globally and by this thread, respectively.
+ let (concatenated, _bytes_allocated_globally, bytes_allocated_by_this_thread) =
+ memory_use(|| {
+ let mut concatenated = concatenate(num_concats, list_array.clone());
+ concatenated.shrink_to_fit(); // This is what we're testing!
Review Comment:
Without this call, this test fails (as it should).
In other words, this test is a regression-test for `shrink_to_fit`
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2500750807
Btw, this is not a breaking change, so this could be added to an `arrow` `0.53.x` patch release (which I would greatly appreciate π)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1858448486
##########
arrow-buffer/src/bytes.rs:
##########
@@ -96,6 +96,30 @@ impl Bytes {
}
}
+ /// Try to reallocate the underlying memory region to a new size (smaller or larger).
+ ///
+ /// Only works for bytes allocated with the standard allocator.
+ /// Returns `Err` if the memory was allocated with a custom allocator,
+ /// or the call to `realloc` failed, for whatever reason.
+ /// In case of `Err`, the [`Bytes`] will remain as it was (i.e. have the old size).
+ pub fn try_realloc(&mut self, new_len: usize) -> Result<(), ()> {
Review Comment:
I've documented why it is safe now, and while doing so found a bug where `new_len = 0` would produce UB. That is now fixed :)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1856881778
##########
arrow-array/src/array/mod.rs:
##########
@@ -365,6 +368,16 @@ impl Array for ArrayRef {
self.as_ref().is_empty()
}
+ fn shrink_to_fit(&mut self) {
+ if let Some(slf) = Arc::get_mut(self) {
+ slf.shrink_to_fit();
+ } else {
+ // TODO(emilk): clone the contents and shrink that.
+ // This can be accomplished if we add `trait Array { fn clone(&self) -> Box<Array>>; }`.
+ // Or we clone using `let clone = self.slice(0, self.len());` and hope that the returned `ArrayRef` is not shared.
Review Comment:
I think any of the following three approaches makes sense here:
## Do nothing
The current approach. If the `ArrayRef` is shared, then we cannot free up memory.
In many ways this makes sense, and I'd be fine with it. We can also start here and improve it later.
## Clone using `slice`
```rs
let mut clone = self.slice(0, self.len());
if let Some(slf) = Arc::get_mut(&mut clone) {
clone.shrink_to_fit();
} else {
// This would be highly unlikely, and fine
}
clone
```
## Clone using a new trait method
Add
``` rs
trait Array {
fn clone(&self) -> Box<dyn Array>;
}
```
(and implement it everywhere)
and then call
```rs
let mut clone = Array::clone(self);
clone.shrink_to_fit();
Arc::new(clone)
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold merged PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1856877614
##########
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
Review Comment:
Ooops⦠my editor is set to trim trailing whitspace on save. Let me know if you want me to reverse.
##########
arrow-array/src/array/mod.rs:
##########
@@ -365,6 +368,16 @@ impl Array for ArrayRef {
self.as_ref().is_empty()
}
+ fn shrink_to_fit(&mut self) {
+ if let Some(slf) = Arc::get_mut(self) {
+ slf.shrink_to_fit();
+ } else {
+ // TODO(emilk): clone the contents and shrink that.
+ // This can be accomplished if we add `trait Array { fn clone(&self) -> Box<Array>>; }`.
+ // Or we clone using `let clone = self.slice(0, self.len());` and hope that the returned `ArrayRef` is not shared.
Review Comment:
I think any of the following three approaches makes sense here:
## Do nothing
The current approach. If the `ArrayRef` is shared, then we cannot free up memory.
## Clone using `slice`
```rs
let mut clone = self.slice(0, self.len());
if let Some(slf) = Arc::get_mut(clone) {
clone.shrink_to_fit();
} else {
// This would be highly unlikely
}
clone
```
## Clone using a new trait method
Add
``` rs
trait Array {
fn clone(&self) -> Box<dyn Array>;
}
(and implement it everywhere)
```
and then call
```rs
let mut clone = Array::clone(self);
clone.shrink_to_fit();
Arc::new(clone)
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1856881778
##########
arrow-array/src/array/mod.rs:
##########
@@ -365,6 +368,16 @@ impl Array for ArrayRef {
self.as_ref().is_empty()
}
+ fn shrink_to_fit(&mut self) {
+ if let Some(slf) = Arc::get_mut(self) {
+ slf.shrink_to_fit();
+ } else {
+ // TODO(emilk): clone the contents and shrink that.
+ // This can be accomplished if we add `trait Array { fn clone(&self) -> Box<Array>>; }`.
+ // Or we clone using `let clone = self.slice(0, self.len());` and hope that the returned `ArrayRef` is not shared.
Review Comment:
I think any of the following three approaches makes sense here:
## Do nothing
The current approach. If the `ArrayRef` is shared, then we cannot free up memory.
In many ways this makes sense, and I'd be fine with it. We can also start here and improve it later.
## Clone using `slice`
```rs
let mut clone = self.slice(0, self.len());
if let Some(slf) = Arc::get_mut(clone) {
clone.shrink_to_fit();
} else {
// This would be highly unlikely
}
clone
```
## Clone using a new trait method
Add
``` rs
trait Array {
fn clone(&self) -> Box<dyn Array>;
}
(and implement it everywhere)
```
and then call
```rs
let mut clone = Array::clone(self);
clone.shrink_to_fit();
Arc::new(clone)
```
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1858775436
##########
arrow-buffer/src/buffer/immutable.rs:
##########
@@ -576,30 +591,30 @@ mod tests {
#[test]
fn test_shrink_to_fit() {
- let original = Buffer::from(&[1, 2, 3, 4, 5, 6, 7]);
- assert_eq!(original.len(), 7);
+ 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(3);
+ 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.len(), 4);
+ assert_eq!(slice.as_slice(), &[2, 3, 4]);
assert_eq!(slice.capacity(), 64);
let mut shrunk = slice;
shrunk.shrink_to_fit();
- assert_eq!(shrunk.len(), 4);
- assert_eq!(shrunk.capacity(), 4);
+ 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(4);
+ 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.len(), 0);
- assert_eq!(empty_slice.capacity(), 4);
+ 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.len(), 0);
- assert_eq!(shrunk_empty.capacity(), 1); // `Buffer` and `Bytes` doesn't support 0-capacity, so we shrink to 1
+ assert_eq!(shrunk_empty.as_slice(), &[]);
+ assert_eq!(shrunk_empty.capacity(), 1); // NOTE: `Buffer` and `Bytes` doesn't support 0-capacity
Review Comment:
They should do, but IIRC you need to use a dangling ptr, there should be some examples of this...
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on code in PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#discussion_r1859090032
##########
arrow-buffer/src/buffer/immutable.rs:
##########
@@ -576,30 +591,30 @@ mod tests {
#[test]
fn test_shrink_to_fit() {
- let original = Buffer::from(&[1, 2, 3, 4, 5, 6, 7]);
- assert_eq!(original.len(), 7);
+ 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(3);
+ 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.len(), 4);
+ assert_eq!(slice.as_slice(), &[2, 3, 4]);
assert_eq!(slice.capacity(), 64);
let mut shrunk = slice;
shrunk.shrink_to_fit();
- assert_eq!(shrunk.len(), 4);
- assert_eq!(shrunk.capacity(), 4);
+ 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(4);
+ 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.len(), 0);
- assert_eq!(empty_slice.capacity(), 4);
+ 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.len(), 0);
- assert_eq!(shrunk_empty.capacity(), 1); // `Buffer` and `Bytes` doesn't support 0-capacity, so we shrink to 1
+ assert_eq!(shrunk_empty.as_slice(), &[]);
+ assert_eq!(shrunk_empty.capacity(), 1); // NOTE: `Buffer` and `Bytes` doesn't support 0-capacity
Review Comment:
For `MutableBuffer` there is special handling for the size=0 case, with a [`dangling_ptr`](https://github.com/rerun-io/arrow-rs/blob/3b7cdebbe306b1e58f9ecfb4ce3b7a1b635a5d07/arrow-buffer/src/buffer/mutable.rs#L471-L472) helper. We could copy all that logic to `Bytes`, but I rather not add all that complexity in this PR.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2500106102
What do you think about adding a `struct ShrinkPolicy` parameter to `shrink_to_fit`? That would allow us to add custom options in the future without breaking backwards compatibility. We could start with
```rust
#[non_exhaustive]
#[derive(Debug, Default)]
pub struct ShrinkPolicy {}
```
Options I could see us adding in the future:
* What to do with shared `ArrayRef`
* What to do with shared `Buffer`
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "emilk (via GitHub)" <gi...@apache.org>.
emilk commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2501194504
MIRI caught a couple of bugs in my unsafe code. I pushed some fixes and test improvements.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `Array::shrink_to_fit(&mut self)` [arrow-rs]
Posted by "tustvold (via GitHub)" <gi...@apache.org>.
tustvold commented on PR #6790:
URL: https://github.com/apache/arrow-rs/pull/6790#issuecomment-2500155885
> What do you think about adding a struct ShrinkPolicy parameter to shrink_to_fit? That would allow us to add custom options in the future without breaking backwards compatibility.
It is an idea with merit, but I think let's keep shrink_to_fit simple. More complex logic to recompute minimal array representations belongs in arrow-select, we want to keep arrow-array as lightweight as possible.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org