add spilling for single mode aggregation

Author: 2010YOUY01

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

@@ -82,6 +82,10 @@ pub(in crate::aggregates) struct AggregateHashTable<AggrMode> {
     /// Output schema: group columns followed by aggregate state or final values.
     pub(super) output_schema: SchemaRef,
 
+    /// Intermediate-state schema used when memory pressure requires the table
+    /// to spill its current state.
+    pub(super) state_schema: SchemaRef,
+
     /// Maximum rows per emitted output batch, from config `batch_size`.
     pub(super) batch_size: usize,
 
@@ -97,6 +101,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         agg: &AggregateExec,
         partition: usize,
         output_schema: SchemaRef,
+        state_schema: SchemaRef,
         batch_size: usize,
         filters: Vec<Option<Arc<dyn PhysicalExpr>>>,
     ) -> Result<Self> {
@@ -133,6 +138,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
             group_by_metrics: GroupByMetrics::new(&agg.metrics, partition),
             input_schema,
             output_schema,
+            state_schema,
             batch_size,
             state: AggregateHashTableState::Building(AggregateHashTableBuffer {
                 group_by: Arc::clone(&agg.group_by),
@@ -282,11 +288,48 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
         }
     }
 
+    pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics {
+        self.group_by_metrics.clone()
+    }
+
     /// Returns the number of distinct groups accumulated so far.
     pub(in crate::aggregates) fn building_group_count(&self) -> usize {
         self.state.building().group_values.len()
     }
 
+    /// Takes every intermediate aggregate state and resets the table so it can
+    /// continue accumulating raw input.
+    ///
+    /// Unlike normal single aggregation output, this materializes intermediate
+    /// states rather than final values. The states can therefore be merged after
+    /// spilling without finalizing the same group more than once.
+    pub(in crate::aggregates) fn take_state_batch(
+        &mut self,
+    ) -> Result<Option<RecordBatch>> {
+        let state_schema = Arc::clone(&self.state_schema);
+        let state = self.state.building_mut();
+        if state.group_values.is_empty() {
+            return Ok(None);
+        }
+
+        let mut output = state.group_values.emit(EmitTo::All)?;
+        for acc in &mut state.accumulators {
+            output.extend(acc.state(EmitTo::All)?);
+        }
+
+        let batch = RecordBatch::try_new(state_schema, output)?;
+        debug_assert!(batch.num_rows() > 0);
+
+        // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
+        // key/index buffers too so the memory reservation can be released
+        // before the batch is sorted for spilling.
+        state.group_values.clear_shrink(0);
+        state.batch_group_indices.clear();
+        state.batch_group_indices.shrink_to_fit();
+
+        Ok(Some(batch))
+    }
+
     pub(in crate::aggregates) fn is_building(&self) -> bool {
         matches!(self.state, AggregateHashTableState::Building(_))
     }

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

@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::Arc;
+
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
@@ -41,6 +43,7 @@ impl AggregateHashTable<FinalMarker> {
             agg,
             partition,
             output_schema,
+            Arc::clone(&agg.input().schema()),
             batch_size,
             vec![None; agg.aggr_expr.len()],
         )

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

@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::Arc;
+
 use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
@@ -34,6 +36,7 @@ impl AggregateHashTable<PartialReduceMarker> {
         Self::new_with_filters(
             agg,
             partition,
+            Arc::clone(&output_schema),
             output_schema,
             batch_size,
             vec![None; agg.aggr_expr.len()],

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

@@ -50,6 +50,7 @@ impl AggregateHashTable<PartialMarker> {
         Self::new_with_filters(
             agg,
             partition,
+            Arc::clone(&output_schema),
             output_schema,
             batch_size,
             agg.filter_expr.iter().cloned().collect(),
@@ -87,6 +88,7 @@ impl AggregateHashTable<PartialMarker> {
             group_by_metrics: self.group_by_metrics.clone(),
             input_schema: Arc::clone(&self.input_schema),
             output_schema: Arc::clone(&self.output_schema),
+            state_schema: Arc::clone(&self.state_schema),
             batch_size: self.batch_size,
             state: AggregateHashTableState::Building(AggregateHashTableBuffer {
                 group_by: Arc::clone(&state.group_by),

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

@@ -35,12 +35,14 @@ impl AggregateHashTable<SingleMarker> {
         agg: &AggregateExec,
         partition: usize,
         output_schema: SchemaRef,
+        state_schema: SchemaRef,
         batch_size: usize,
     ) -> Result<Self> {
         Self::new_with_filters(
             agg,
             partition,
             output_schema,
+            state_schema,
             batch_size,
             agg.filter_expr.iter().cloned().collect(),
         )

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

@@ -1274,12 +1274,7 @@ impl AggregateExec {
             && self.group_by.is_single()
     }
 
-    fn should_use_single_hash_stream(&self, context: &TaskContext) -> bool {
-        // TODO: implement memory-limited path and remove this limitation
-        if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) {
-            return false;
-        }
-
+    fn should_use_single_hash_stream(&self, _context: &TaskContext) -> bool {
         matches!(
             self.mode,
             AggregateMode::Single | AggregateMode::SinglePartitioned
@@ -3767,7 +3762,7 @@ mod tests {
         let aggregates_v0: Vec<Arc<AggregateFunctionExpr>> =
             vec![Arc::new(test_median_agg_expr(Arc::clone(&input_schema))?)];
 
-        // use fast-path in `grouped_hash_stream.rs`.
+        // Use the fast path in `single_stream.rs`.
         let aggregates_v2: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new(
             AggregateExprBuilder::new(avg_udaf(), vec![col("b", &input_schema)?])
                 .schema(Arc::clone(&input_schema))
@@ -3800,7 +3795,7 @@ mod tests {
                     assert!(matches!(stream, StreamType::GroupedHash(_)));
                 }
                 2 => {
-                    assert!(matches!(stream, StreamType::GroupedHash(_)));
+                    assert!(matches!(stream, StreamType::SingleHash(_)));
                 }
                 _ => panic!("Unknown version: {version}"),
             }
@@ -4135,15 +4130,14 @@ mod tests {
         Ok(())
     }
 
-    /// Spilling behavior is not implemented for single hash stream yet, so fall
-    /// back to the existing `GroupedHashAggregateStream`.
+    /// Single hash aggregation supports finite memory.
     #[tokio::test]
     async fn single_aggregate_with_memory_limit_planning() -> Result<()> {
         let single = single_test_aggregate()?;
         let task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?;
 
         let stream = single.execute_typed(0, &task_ctx)?;
-        assert!(matches!(stream, StreamType::GroupedHash(_)));
+        assert!(matches!(stream, StreamType::SingleHash(_)));
 
         Ok(())
     }
@@ -5664,9 +5658,11 @@ mod tests {
             Field::new("b", DataType::Float64, false),
         ]));
 
+        let group_keys = [2, 3, 4, 4].repeat(1_000);
+        let values = [1.0, 2.0, 3.0, 4.0].repeat(1_000);
         let batches = vec![
-            create_record_batch(&schema, (vec![2, 3, 4, 4], vec![1.0, 2.0, 3.0, 4.0]))?,
-            create_record_batch(&schema, (vec![2, 3, 4, 4], vec![1.0, 2.0, 3.0, 4.0]))?,
+            create_record_batch(&schema, (group_keys.clone(), values.clone()))?,
+            create_record_batch(&schema, (group_keys, values))?,
         ];
         let plan: Arc<dyn ExecutionPlan> =
             TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?;
@@ -5766,9 +5762,9 @@ mod tests {
     #[tokio::test]
     async fn test_aggregate_with_spill_if_necessary() -> Result<()> {
         // test with spill
-        run_test_with_spill_pool_if_necessary(2_000, true).await?;
+        run_test_with_spill_pool_if_necessary(20_000, true).await?;
         // test without spill
-        run_test_with_spill_pool_if_necessary(20_000, false).await?;
+        run_test_with_spill_pool_if_necessary(200_000, false).await?;
         Ok(())
     }
 
@@ -6744,11 +6740,6 @@ mod tests {
                     matches!(root, DataFusionError::ResourcesExhausted(_)),
                     "Expected ResourcesExhausted, got: {root}",
                 );
-                let msg = root.to_string();
-                assert!(
-                    msg.contains("Failed to reserve memory for sort during spill"),
-                    "Expected sort reservation error, got: {msg}",
-                );
             }
         }
 

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

@@ -97,6 +97,9 @@ struct OrderedFinalSpillContext {
 enum OrderedFinalAggregateState {
     ReadingInput {
         table: OrderedAggregateTable<FinalMarker>,
+        /// None if either
+        /// - Disk Manager doesn't enable temporary file creation
+        /// - The group keys are fully ordered, it's expected to use bounded memory
         spill_context: Option<Box<OrderedFinalSpillContext>>,
     },
     Spilling {
@@ -292,7 +295,7 @@ impl OrderedFinalAggregateStream {
         clippy::too_many_arguments,
         reason = "keeps replay metric reuse explicit"
     )]
-    fn new_with_input_and_metrics(
+    pub(in crate::aggregates) fn new_with_input_and_metrics(
         agg: &AggregateExec,
         context: &Arc<TaskContext>,
         partition: usize,
@@ -440,6 +443,8 @@ impl OrderedFinalAggregateStream {
                     Ok(()) => {}
                     Err(e @ DataFusionError::ResourcesExhausted(_)) => {
                         let Some(spill_context) = spill_context else {
+                            // `None` means spilling is not supported, see comments
+                            // at `OrderedFinalAggregateState` for details.
                             return ControlFlow::Break((
                                 Poll::Ready(Some(Err(e))),
                                 OrderedFinalAggregateState::Done,