diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index 9604087c33334..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 { @@ -280,14 +289,25 @@ 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.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); + } } } @@ -655,4 +675,38 @@ mod tests { 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"); + + // 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(), + 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..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 @@ -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,28 @@ 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); + + // 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 23ea4e7ed3f88..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 @@ -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,29 @@ 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); + + // 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); + } } 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(()) + } }