From 141622d6d38fc45d0858d07aac5481e99795577b Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:30:18 +0300 Subject: [PATCH 1/3] fix: clear_shrink(0) on rows should reset to original size + single group by fixes in clear_shrink the clear_shrink was only called after emit, so this is why the bug was not shown before --- .../src/aggregates/group_values/row.rs | 36 ++++++++++++++++--- .../group_values/single_group_by/bytes.rs | 16 +++++++++ .../single_group_by/bytes_view.rs | 18 ++++++++++ .../group_values/single_group_by/primitive.rs | 31 ++++++++++++++++ 4 files changed, 97 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 01e9f3eaa71ee..075e93bb6407d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -283,10 +283,15 @@ impl GroupValues for GroupValuesRows { } fn clear_shrink(&mut self, num_rows: usize) { - self.group_values = self.group_values.take().map(|mut rows| { - rows.clear(); - rows - }); + self.group_values = if num_rows == 0 { + None + } else { + self.group_values.take().map(|mut rows| { + // This only clear and do not shrink the capacity + rows.clear(); + rows + }) + }; self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); @@ -484,4 +489,27 @@ mod tests { assert_eq!(groups, vec![3]); Ok(()) } + + #[test] + fn clear_shrink_0_after_intern_should_return_to_original_size() { + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let schema = Arc::new(Schema::new(vec![Field::new( + "group", + DataType::List(field), + true, + )])); + let mut group_values = GroupValuesRows::try_new(schema).unwrap(); + let initial_size = group_values.size(); + let input = Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3)]), + Some(vec![Some(1), Some(2)]), + ])) as ArrayRef; + let mut groups = vec![]; + group_values.intern(&[input], &mut groups).unwrap(); + assert_ne!(group_values.size(), initial_size, "should save some data"); + group_values.clear_shrink(0); + assert_eq!(group_values.size(), initial_size, "should release memory back to original size"); + } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 9f7b4b4e91cba..4c7c75414d656 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -141,6 +141,7 @@ impl GroupValues for GroupValuesBytes { // release the map's allocations rather than restoring the warm up // capacities that `take` keeps for the emit path. self.map.clear_and_release(); + self.num_groups = 0; } } @@ -200,4 +201,19 @@ mod tests { group_values.emit(EmitTo::All).unwrap(); assert!(group_values.size() > empty + INITIAL_BUFFER_CAPACITY); } + + #[test] + fn clear_shrink_should_reset_len() { + let mut group_values = GroupValuesBytes::::new(OutputType::Utf8); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..10).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + assert_eq!(group_values.len(), 10); + group_values.clear_shrink(0); + assert_eq!(group_values.len(), 0); + } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 23ea4e7ed3f88..a0c2896b494e3 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -143,6 +143,7 @@ impl GroupValues for GroupValuesBytesView { // release the map's allocations rather than restoring the warm up // capacity that `take` keeps for the emit path. self.map.clear_and_release(); + self.num_groups = 0; } } @@ -201,4 +202,21 @@ mod tests { group_values.emit(EmitTo::All).unwrap(); assert!(group_values.size() > empty + INITIAL_MAP_CAPACITY); } + + + #[test] + fn clear_shrink_should_reset_len() { + let mut group_values = GroupValuesBytesView::new(OutputType::Utf8View); + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..10).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + assert_eq!(group_values.len(), 10); + + group_values.clear_shrink(0); + assert_eq!(group_values.len(), 0); + } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index 21b62457e3831..9aba200563155 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -273,6 +273,9 @@ where self.values.shrink_to(num_rows); self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + + // Reset the null group index + self.null_group = None; } } @@ -322,4 +325,32 @@ mod tests { Ok(()) } + + #[test] + fn clear_shrink_reset_null_group() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + + // Intern some values including a null + let arr: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(2)])); + let mut groups = vec![]; + gv.intern(&[arr], &mut groups)?; + + assert_eq!(groups.len(), 3); + + let null_group = groups[1]; + assert_eq!(null_group, 1); + + // Clear and shrink + gv.clear_shrink(0); + + let arr: ArrayRef = Arc::new(Int32Array::from(vec![None::])); + let mut groups = vec![]; + gv.intern(&[arr], &mut groups)?; + assert_eq!(groups.len(), 1); + + let new_null_group = groups[0]; + assert_eq!(new_null_group, 0); + + Ok(()) + } } From 35d48e009056e1852346b9ed50265a84412ce697 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:34:58 +0300 Subject: [PATCH 2/3] format --- datafusion/physical-plan/src/aggregates/group_values/row.rs | 6 +++++- .../aggregates/group_values/single_group_by/bytes_view.rs | 1 - 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 075e93bb6407d..9e7dbbdf63adc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -510,6 +510,10 @@ mod tests { group_values.intern(&[input], &mut groups).unwrap(); assert_ne!(group_values.size(), initial_size, "should save some data"); group_values.clear_shrink(0); - assert_eq!(group_values.size(), initial_size, "should release memory back to original size"); + assert_eq!( + group_values.size(), + initial_size, + "should release memory back to original size" + ); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index a0c2896b494e3..b9201a2b3527c 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -203,7 +203,6 @@ mod tests { assert!(group_values.size() > empty + INITIAL_MAP_CAPACITY); } - #[test] fn clear_shrink_should_reset_len() { let mut group_values = GroupValuesBytesView::new(OutputType::Utf8View); From 760ef21a155e9e812a401327b4da0bdde6cc47b1 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 4 Oct 2026 13:21:58 +0300 Subject: [PATCH 3/3] code review changes --- .../src/aggregates/group_values/row.rs | 32 ++++++++++++++++--- .../group_values/single_group_by/bytes.rs | 9 ++++++ .../single_group_by/bytes_view.rs | 9 ++++++ 3 files changed, 45 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index ddd664caa231f..8a73ec6e50966 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -75,6 +75,9 @@ pub struct GroupValuesRows { /// reused buffer to store rows rows_buffer: Rows, + /// The initial `rows_buffer` size, used to determine if it should be recreated in `clear_shrink`. + initial_rows_buffer_size: usize, + /// Random state for creating hashes random_state: RandomState, } @@ -94,11 +97,8 @@ impl GroupValuesRows { let map = HashTable::with_capacity(0); - let starting_rows_capacity = 1000; - - let starting_data_capacity = 64 * starting_rows_capacity; - let rows_buffer = - row_converter.empty_rows(starting_rows_capacity, starting_data_capacity); + let rows_buffer = Self::new_rows_buffer(&row_converter); + let initial_rows_buffer_size = rows_buffer.size(); Ok(Self { schema, row_converter, @@ -106,9 +106,18 @@ impl GroupValuesRows { group_values: None, hashes_buffer: Default::default(), rows_buffer, + initial_rows_buffer_size, random_state: crate::aggregates::AGGREGATION_HASH_SEED, }) } + + /// Create the reused `rows_buffer` with its initial capacities + fn new_rows_buffer(row_converter: &RowConverter) -> Rows { + let starting_rows_capacity = 1000; + + let starting_data_capacity = 64 * starting_rows_capacity; + row_converter.empty_rows(starting_rows_capacity, starting_data_capacity) + } } impl GroupValues for GroupValuesRows { @@ -293,6 +302,12 @@ impl GroupValues for GroupValuesRows { self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared self.hashes_buffer.clear(); self.hashes_buffer.shrink_to(num_rows); + + if num_rows == 0 && self.rows_buffer.size() > self.initial_rows_buffer_size { + // `rows_buffer` only gets cleared between batches, so it keeps any + // capacity it grew to; recreate it to release that memory + self.rows_buffer = Self::new_rows_buffer(&self.row_converter); + } } } @@ -680,6 +695,13 @@ mod tests { let mut groups = vec![]; group_values.intern(&[input], &mut groups).unwrap(); assert_ne!(group_values.size(), initial_size, "should save some data"); + + // A large batch so any reused buffers grow as well + let input = Arc::new(ListArray::from_iter_primitive::( + (0..10_000).map(|i| Some((0..20).map(move |j| Some(i * 20 + j)))), + )) as ArrayRef; + group_values.intern(&[input], &mut groups).unwrap(); + group_values.clear_shrink(0); assert_eq!( group_values.size(), diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 4c7c75414d656..3c92ce787dc2b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -215,5 +215,14 @@ mod tests { assert_eq!(group_values.len(), 10); group_values.clear_shrink(0); assert_eq!(group_values.len(), 0); + + // The first new distinct value after clearing should get group index 0 + let values: ArrayRef = Arc::new(StringArray::from_iter_values([ + "new value", + "another value", + ])); + group_values.intern(&[values], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1]); + assert_eq!(group_values.len(), 2); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index b9201a2b3527c..a2e3573ce889f 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -217,5 +217,14 @@ mod tests { group_values.clear_shrink(0); assert_eq!(group_values.len(), 0); + + // The first new distinct value after clearing should get group index 0 + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values([ + "new value", + "another value", + ])); + group_values.intern(&[values], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1]); + assert_eq!(group_values.len(), 2); } }