From 7380a865572d4f94241a9dfe1240d0faae499959 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 9 Oct 2026 16:45:47 +0800 Subject: [PATCH 1/2] Implement min_max_bytes with capacity tracking and regression tests - Counts retained outer/inner capacities; size() stays O(1). - Adds 3 regressions: capacity retention, growth, emit/reuse. --- .../src/min_max/min_max_bytes.rs | 120 +++++++++++++++++- 1 file changed, 116 insertions(+), 4 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 7d26143eb5790..f236a35f43608 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>>, /// 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(); } } } @@ -516,6 +522,7 @@ impl MinMaxBytesState { fn emit_to(&mut self, emit_to: EmitTo) -> (usize, Vec>>) { 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), @@ -528,13 +535,20 @@ impl MinMaxBytesState { .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) .sum(); self.total_data_bytes -= first_data_capacity; + self.total_data_capacity -= first_min_maxes + .iter() + .flatten() + .map(Vec::capacity) + .sum::(); (first_data_capacity, first_min_maxes) } } } fn size(&self) -> usize { - self.total_data_bytes + self.min_max.len() * size_of::>>() + // Only retained heap allocations: the outer slots already include the + // inline Option> descriptors. + self.total_data_capacity + self.min_max.capacity() * size_of::>>() } } @@ -543,6 +557,104 @@ mod tests { use super::*; use arrow::array::StringArray; + #[test] + fn size_counts_outer_capacity() { + let mut state = MinMaxBytesState::new(DataType::Utf8); + assert_eq!(state.size(), 0); + + state.min_max = Vec::with_capacity(16); + state.min_max.resize(3, None); + assert_eq!( + state.size(), + state.min_max.capacity() * size_of::>>() + ); + } + + #[test] + fn size_retains_inner_capacity() -> Result<()> { + for is_min in [true, false] { + let mut accumulator = if is_min { + MinMaxBytesAccumulator::new_min(DataType::Utf8) + } else { + MinMaxBytesAccumulator::new_max(DataType::Utf8) + }; + let long = if is_min { + "z".repeat(128) + } else { + "a".repeat(128) + }; + let short = if is_min { "a" } else { "z" }; + let values: ArrayRef = Arc::new(StringArray::from(vec![long.as_str(), "m"])); + accumulator.update_batch(&[values], &[0, 1], None, 2)?; + let outer_bytes = + accumulator.inner.min_max.capacity() * size_of::>>(); + let inner_bytes: usize = accumulator + .inner + .min_max + .iter() + .flatten() + .map(Vec::capacity) + .sum(); + assert_eq!(accumulator.size(), outer_bytes + inner_bytes); + + let values: ArrayRef = Arc::new(StringArray::from(vec![short])); + accumulator.update_batch(&[values], &[0], None, 2)?; + assert_eq!(accumulator.inner.total_data_bytes, short.len() + 1); + assert_eq!(accumulator.size(), outer_bytes + inner_bytes); + + let emitted_capacity = + accumulator.inner.min_max[0].as_ref().unwrap().capacity(); + let (data_bytes, emitted) = accumulator.inner.emit_to(EmitTo::First(1)); + assert_eq!(data_bytes, short.len()); + assert_eq!(emitted[0].as_deref(), Some(short.as_bytes())); + assert_eq!( + accumulator.size(), + accumulator.inner.min_max.capacity() * size_of::>>() + + inner_bytes + - emitted_capacity + ); + + let (data_bytes, emitted) = accumulator.inner.emit_to(EmitTo::All); + assert_eq!(data_bytes, 1); + assert_eq!(emitted[0].as_deref(), Some(b"m".as_slice())); + assert_eq!(accumulator.size(), 0); + } + Ok(()) + } + + #[test] + fn size_tracks_growth_and_emit_reuse() -> Result<()> { + let mut state = MinMaxBytesState::new(DataType::Binary); + state.min_max.resize(3, None); + state.set_value(0, b"a"); + state.set_value(0, &[b'z'; 128]); + state.set_value(1, b"b"); + let inner_bytes: usize = state.min_max.iter().flatten().map(Vec::capacity).sum(); + assert_eq!( + state.size(), + state.min_max.capacity() * size_of::>>() + inner_bytes + ); + + // Emit the larger side, which replaces the outer allocation. + let (data_bytes, emitted) = state.emit_to(EmitTo::First(2)); + assert_eq!(data_bytes, 129); + assert_eq!(emitted.len(), 2); + assert_eq!( + state.size(), + state.min_max.capacity() * size_of::>>() + ); + + state.emit_to(EmitTo::All); + assert_eq!(state.size(), 0); + state.update_batch([Some(b"x".as_slice())], &[0], 1, |a, b| a < b)?; + assert_eq!( + state.size(), + state.min_max.capacity() * size_of::>>() + + state.min_max[0].as_ref().unwrap().capacity() + ); + Ok(()) + } + #[test] fn preserving_selected_min_values() -> Result<()> { let mut accumulator = MinMaxBytesAccumulator::new_min(DataType::Utf8); From e7a07f48dfa58b96a0f07c4cb72672bcc8587bf1 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 9 Oct 2026 16:58:54 +0800 Subject: [PATCH 2/2] Fix regression in min_max_bytes output validation - Existing regression now checks evaluate(EmitTo::First(1)) output + remaining size, both min/max. - Kept direct full-emission byte-length checks. --- datafusion/functions-aggregate/src/min_max/min_max_bytes.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 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 f236a35f43608..e071a999d6d03 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -604,9 +604,8 @@ mod tests { let emitted_capacity = accumulator.inner.min_max[0].as_ref().unwrap().capacity(); - let (data_bytes, emitted) = accumulator.inner.emit_to(EmitTo::First(1)); - assert_eq!(data_bytes, short.len()); - assert_eq!(emitted[0].as_deref(), Some(short.as_bytes())); + let emitted = accumulator.evaluate(EmitTo::First(1))?; + assert_eq!(emitted.as_string::(), &StringArray::from(vec![short])); assert_eq!( accumulator.size(), accumulator.inner.min_max.capacity() * size_of::>>()