fix: clear_shrink(0) on rows should reset to original size + single group by fixes in clear_shrink

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,21 +97,27 @@ 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,
             map,
             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::<Int32Type, _, _>(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::<Int32Type, _, _>(
+            (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"
+        );
+    }
 }

datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs

@@ -141,6 +141,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesBytes<O> {
         // 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::<i32>::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);
+    }
 }

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);
+    }
 }

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::<Int32Type>::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::<i32>]));
+        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(())
+    }
 }