ariel-miculas commented on code in PR #22416:
URL: https://github.com/apache/datafusion/pull/22416#discussion_r3285982930
##########
datafusion/common/src/utils/mod.rs:
##########
@@ -395,6 +395,74 @@ pub fn longest_consecutive_prefix<T: Borrow<usize>>(
count
}
+/// Splits `vec` at index `n`, returning the first `n` elements and leaving the
+/// remaining `vec.len() - n` elements in `vec`.
+///
+/// Allocates for whichever side is smaller, so the new allocation is
+/// `min(n, vec.len() - n)` rather than always `n` (as
`vec.drain(0..n).collect()`
+/// would). This matters when the split emits a prefix under memory pressure,
+/// where `n` can be close to `vec.len()`.
+pub fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize) -> Vec<T> {
+ if n * 2 <= vec.len() {
+ vec.drain(0..n).collect()
+ } else {
+ let remaining = vec.split_off(n);
+ // The emitted prefix keeps the original (larger) allocation.
datafusion
+ // accounts memory by capacity rather than length, so shrink it to
avoid
+ // handing off a short Vec that still reserves the pre-split capacity.
+ vec.shrink_to_fit();
Review Comment:
This was my bad for suggesting this, turns out it's a more nuanced change
that needs more tests. In
https://github.com/apache/datafusion/blob/97d60af9feea027efeeb911bdd195bf2ed704ee4/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs#L384-L391
there's
```
first_n_offsets.push(offset_n);
```
and if we're shrinking the vec to the minimum capacity, pushing to it will
cause reallocation and a greater capacity than the original one.
It's probably best to remove the shrink_to_fit for 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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
