chore(PartialReduceHashAggregateStream): convert to async generators and cleanup

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -410,10 +410,6 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         matches!(self.state, AggregateHashTableState::Building(_))
     }
 
-    pub(in crate::aggregates) fn is_done(&self) -> bool {
-        matches!(self.state, AggregateHashTableState::Done)
-    }
-
     pub(super) fn start_outputting(&mut self) {
         let AggregateHashTableState::Building(mut state) =
             std::mem::replace(&mut self.state, AggregateHashTableState::Done)

datafusion/physical-plan/src/aggregates/mod.rs

@@ -716,7 +716,7 @@ impl From<StreamType> for SendableRecordBatchStream {
         match stream {
             StreamType::AggregateStream(stream) => Box::pin(stream),
             StreamType::PartialHash(stream) => stream.into_stream(),
-            StreamType::PartialReduceHash(stream) => Box::pin(stream),
+            StreamType::PartialReduceHash(stream) => stream.into_stream(),
             StreamType::FinalHash(stream) => stream.into_stream(),
             StreamType::SingleHash(stream) => stream.into_stream(),
             StreamType::OrderedPartialAggregate(stream) => stream.into_stream(),
@@ -6069,9 +6069,11 @@ mod tests {
         let stream: SendableRecordBatchStream = stream.into();
         let output = collect(stream).await?;
 
-        // Same flush cadence as the column-backed test: one flush per input
-        // batch, each sliced into batches of 2 and 1 rows.
-        assert_eq!(output.len(), 2 * num_input_batches);
+        // One flush per input batch. The memory limit only fits the table, so
+        // holding the materialized flush while slicing it does not fit and each
+        // flush is emitted whole instead of in `batch_size` slices.
+        assert_eq!(output.len(), num_input_batches);
+        assert!(output.iter().all(|batch| batch.num_rows() == 3));
         assert_snapshot!(batches_to_string(&output), @r"
         +---+---+-------------+
         | a | n | SUM(b)[sum] |