You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/14 05:35:52 UTC
(datafusion) 01/01: Fix retained-capacity accounting for grouped byte MIN/MAX (#26154)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26154-92323f7f146bb43931193bf005f35804dbca15b7
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0d79252f61bd8777990b2749c216ffd07079c30d
Author: kosiew <ko...@gmail.com>
AuthorDate: Wed Oct 14 05:35:27 2026 +0000
Fix retained-capacity accounting for grouped byte MIN/MAX (#26154)
## Which issue does this PR close?
- Part of #23393
## Rationale for this change
Grouped byte MIN/MAX accumulators can underreport memory usage because
their accounting is based on the current data length and number of group
slots rather than the memory allocations they retain.
For example, replacing a 128-byte string with a one-byte string reuses
the existing vector allocation, but the previous accounting reports only
the new string length. Similarly, unused capacity in the outer group
vector is not counted.
This change makes memory accounting reflect retained allocations,
providing more accurate memory usage estimates for grouped MIN/MAX
operations.
## What changes are included in this PR?
- Add `total_data_capacity` to `MinMaxBytesState` to track retained
inner `Vec<u8>` capacities independently of `total_data_bytes`, which
remains responsible for output preallocation.
- Update `set_value()` to maintain capacity accounting when values are
inserted, replaced, or cause an existing vector to grow.
- Update `size()` to account for the outer `min_max` vector's capacity
and the cached inner vector capacities, without double-counting the
inline `Option<Vec<u8>>` descriptors.
- Keep `size()` O(1) by maintaining the inner capacity total
incrementally.
- Update `emit_to()` to reset capacity accounting on full emission and
subtract emitted inner vector capacities on partial emission.
## Are these changes tested?
Three regression tests are added:
- `size_counts_outer_capacity`: Verifies that memory accounting includes
allocated outer vector capacity, even when fewer group slots are
initialized.
- `size_retains_inner_capacity`: Tests both MIN and MAX accumulators,
verifying that replacing a long string with a shorter one preserves
retained-capacity accounting and that partial and full emission update
it correctly.
- `size_tracks_growth_and_emit_reuse`: Verifies capacity accounting when
an inner vector grows, groups are partially emitted, and the state is
reused after full emission.
## Are there any user-facing changes?
No changes to SQL MIN/MAX results or public APIs are intended.
The change corrects internal memory accounting for grouped byte MIN/MAX
accumulators, allowing retained allocations to be reported more
accurately.
## LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated
content has been manually reviewed.
---
.../src/min_max/min_max_bytes.rs | 94 +++++++++++++++++++---
1 file changed, 83 insertions(+), 11 deletions(-)
diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs
index 7d26143eb5..7fb30a4f9b 100644
--- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs
+++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs
@@ -416,9 +416,10 @@ struct MinMaxBytesState {
min_max: Vec<Option<Vec<u8>>>,
/// The data type of the array
data_type: DataType,
- /// The total bytes of the string data (for pre-allocating the final array,
- /// and tracking memory usage)
+ /// The total length of the string data, for pre-allocating the final array.
total_data_bytes: usize,
+ /// The retained inner vector capacities, cached to keep `size()` O(1).
+ total_data_capacity: usize,
}
/// Implement the MinMaxBytesAccumulator with a comparison function
@@ -433,6 +434,7 @@ impl MinMaxBytesState {
min_max: vec![],
data_type,
total_data_bytes: 0,
+ total_data_capacity: 0,
}
}
@@ -440,15 +442,19 @@ impl MinMaxBytesState {
fn set_value(&mut self, group_index: usize, new_val: &[u8]) {
match self.min_max[group_index].as_mut() {
None => {
- self.min_max[group_index] = Some(new_val.to_vec());
+ let value = new_val.to_vec();
+ self.total_data_capacity += value.capacity();
+ self.min_max[group_index] = Some(value);
self.total_data_bytes += new_val.len();
}
Some(existing_val) => {
// Copy data over to avoid re-allocating
self.total_data_bytes -= existing_val.len();
self.total_data_bytes += new_val.len();
+ self.total_data_capacity -= existing_val.capacity();
existing_val.clear();
existing_val.extend_from_slice(new_val);
+ self.total_data_capacity += existing_val.capacity();
}
}
}
@@ -509,13 +515,14 @@ impl MinMaxBytesState {
/// Emits the specified min_max values
///
- /// Returns (data_capacity, min_maxes), updating the current value of total_data_bytes
+ /// Returns (data_bytes, min_maxes), updating the current value of total_data_bytes
///
- /// - `data_capacity`: the total length of all strings and their contents,
+ /// - `data_bytes`: the total byte length of the emitted values,
/// - `min_maxes`: the actual min/max values for each group
fn emit_to(&mut self, emit_to: EmitTo) -> (usize, Vec<Option<Vec<u8>>>) {
match emit_to {
EmitTo::All => {
+ self.total_data_capacity = 0;
(
std::mem::take(&mut self.total_data_bytes), // reset total bytes and min_max
std::mem::take(&mut self.min_max),
@@ -523,18 +530,23 @@ impl MinMaxBytesState {
}
EmitTo::First(n) => {
let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n);
- let first_data_capacity: usize = first_min_maxes
+ let (first_data_bytes, first_data_capacity) = first_min_maxes
.iter()
- .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0))
- .sum();
- self.total_data_bytes -= first_data_capacity;
- (first_data_capacity, first_min_maxes)
+ .flatten()
+ .fold((0, 0), |(bytes, capacity), value| {
+ (bytes + value.len(), capacity + value.capacity())
+ });
+ self.total_data_bytes -= first_data_bytes;
+ self.total_data_capacity -= first_data_capacity;
+ (first_data_bytes, first_min_maxes)
}
}
}
fn size(&self) -> usize {
- self.total_data_bytes + self.min_max.len() * size_of::<Option<Vec<u8>>>()
+ // Only retained heap allocations: the outer slots already include the
+ // inline Option<Vec<u8>> descriptors.
+ self.total_data_capacity + self.min_max.capacity() * size_of::<Option<Vec<u8>>>()
}
}
@@ -543,6 +555,66 @@ mod tests {
use super::*;
use arrow::array::StringArray;
+ fn expected_size(state: &MinMaxBytesState) -> usize {
+ state.min_max.capacity() * size_of::<Option<Vec<u8>>>()
+ + state
+ .min_max
+ .iter()
+ .flatten()
+ .map(Vec::capacity)
+ .sum::<usize>()
+ }
+
+ #[test]
+ fn size_tracks_growth_and_emit_reuse() -> Result<()> {
+ let mut state = MinMaxBytesState::new(DataType::Binary);
+ assert_eq!(state.size(), 0);
+
+ state.min_max = Vec::with_capacity(16);
+ assert_eq!(state.size(), expected_size(&state));
+ state.min_max.resize(4, None);
+ assert_eq!(state.size(), expected_size(&state));
+
+ state.set_value(0, b"a");
+ state.set_value(0, &[b'z'; 128]);
+ state.set_value(1, b"b");
+ state.set_value(2, b"c");
+ state.set_value(3, b"d");
+ assert_eq!(state.size(), expected_size(&state));
+ let size_before_replacement = state.size();
+
+ state.set_value(0, b"x");
+ assert_eq!(state.size(), size_before_replacement);
+ assert!(state.min_max[0].as_ref().unwrap().capacity() >= 128);
+ assert_eq!(state.total_data_bytes, 4);
+ assert_eq!(state.size(), expected_size(&state));
+
+ let (bytes, emitted) = state.emit_to(EmitTo::First(1));
+ assert_eq!(emitted[0].as_deref(), Some(b"x".as_slice()));
+ assert_eq!(bytes, 1);
+ assert_eq!(state.size(), expected_size(&state));
+ assert_eq!(state.total_data_bytes, 3);
+
+ let (bytes, emitted) = state.emit_to(EmitTo::First(2));
+ assert_eq!(emitted[0].as_deref(), Some(b"b".as_slice()));
+ assert_eq!(emitted[1].as_deref(), Some(b"c".as_slice()));
+ assert_eq!(bytes, 2);
+ assert_eq!(state.size(), expected_size(&state));
+ assert_eq!(state.total_data_bytes, 1);
+
+ let (bytes, emitted) = state.emit_to(EmitTo::All);
+ assert_eq!(emitted[0].as_deref(), Some(b"d".as_slice()));
+ assert_eq!(bytes, 1);
+ assert_eq!(state.size(), 0);
+ assert_eq!(state.total_data_bytes, 0);
+
+ state.update_batch([Some(b"x".as_slice())], &[0], 1, |a, b| a < b)?;
+ assert_eq!(state.min_max[0].as_deref(), Some(b"x".as_slice()));
+ assert_eq!(state.total_data_bytes, 1);
+ assert_eq!(state.size(), expected_size(&state));
+ Ok(())
+ }
+
#[test]
fn preserving_selected_min_values() -> Result<()> {
let mut accumulator = MinMaxBytesAccumulator::new_min(DataType::Utf8);
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org