refactor: Split hash aggregation logic into separated streams

datafusion/common/src/config.rs

@@ -587,6 +587,17 @@ config_namespace! {
         /// the new schema verification step.
         pub skip_physical_aggregate_schema_check: bool, default = false
 
+        /// Temporary switch for aggregate stream implementations that are being
+        /// migrated from `GroupedHashAggregateStream`.
+        ///
+        /// When set to true, DataFusion tries the migrated implementations when
+        /// their preconditions are satisfied. When set to false, grouped
+        /// aggregation falls back to `GroupedHashAggregateStream`. This option
+        /// will be removed after the migration is finished.
+        ///
+        /// See <https://github.com/apache/datafusion/issues/22710> for details.
+        pub enable_migration_aggregate: bool, default = false
+
         /// Sets the compression codec used when spilling data to disk.
         ///
         /// Since datafusion writes spill files using the Arrow IPC Stream format,

datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs

@@ -41,6 +41,7 @@ use crate::fuzz_cases::aggregation_fuzzer::data_generator::Dataset;
 ///   - `batch_size`
 ///   - `target_partitions`
 ///   - `skip_partial parameters`
+///   - `enable_migration_aggregate`
 ///   - hint `sorted` or not
 ///   - `spilling` or not (TODO, I think a special `MemoryPool` may be needed
 ///     to support this)
@@ -96,11 +97,13 @@ impl SessionContextGenerator {
         let batch_size = self.max_batch_size;
         let target_partitions = 1;
         let skip_partial_params = SkipPartialParams::ensure_not_trigger();
+        let enable_migration_aggregate = false;
 
         let builder = GeneratedSessionContextBuilder {
             batch_size,
             target_partitions,
             skip_partial_params,
+            enable_migration_aggregate,
             sort_hint: false,
             table_name: self.table_name.clone(),
             table_provider: Arc::new(provider),
@@ -120,6 +123,7 @@ impl SessionContextGenerator {
         //   - `batch_size`, from range: [1, `total_rows_num`]
         //   - `target_partitions`, from range: [1, cpu_num]
         //   - `skip_partial`, trigger or not trigger currently for simplicity
+        //   - `enable_migration_aggregate`, true or false
         //   - `sorted`, if found a sorted dataset, will or will not push down this information
         //   - `spilling`(TODO)
         let batch_size = rng.random_range(1..=self.max_batch_size);
@@ -131,6 +135,8 @@ impl SessionContextGenerator {
         let skip_partial_params =
             self.candidate_skip_partial_params[skip_partial_params_idx];
 
+        let enable_migration_aggregate = rng.random_bool(0.5);
+
         let (provider, sort_hint) =
             if rng.random_bool(0.5) && !self.dataset.sort_keys.is_empty() {
                 // Sort keys exist and random to push down
@@ -150,6 +156,7 @@ impl SessionContextGenerator {
             target_partitions,
             sort_hint,
             skip_partial_params,
+            enable_migration_aggregate,
             table_name: self.table_name.clone(),
             table_provider: Arc::new(provider),
         };
@@ -173,6 +180,7 @@ struct GeneratedSessionContextBuilder {
     target_partitions: usize,
     sort_hint: bool,
     skip_partial_params: SkipPartialParams,
+    enable_migration_aggregate: bool,
     table_name: String,
     table_provider: Arc<dyn TableProvider>,
 }
@@ -197,6 +205,10 @@ impl GeneratedSessionContextBuilder {
             "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold",
             &ScalarValue::Float64(Some(self.skip_partial_params.ratio_threshold)),
         );
+        session_config = session_config.set_bool(
+            "datafusion.execution.enable_migration_aggregate",
+            self.enable_migration_aggregate,
+        );
 
         let ctx = SessionContext::new_with_config(session_config);
         ctx.register_table(self.table_name, self.table_provider)?;
@@ -206,6 +218,7 @@ impl GeneratedSessionContextBuilder {
             target_partitions: self.target_partitions,
             sort_hint: self.sort_hint,
             skip_partial_params: self.skip_partial_params,
+            enable_migration_aggregate: self.enable_migration_aggregate,
         };
 
         Ok(SessionContextWithParams { ctx, params })
@@ -220,6 +233,7 @@ pub struct SessionContextParams {
     target_partitions: usize,
     sort_hint: bool,
     skip_partial_params: SkipPartialParams,
+    enable_migration_aggregate: bool,
 }
 
 /// Partial skipping parameters

datafusion/physical-plan/src/aggregates/group_values/metrics.rs

@@ -59,6 +59,7 @@ mod tests {
     use arrow::record_batch::RecordBatch;
     use datafusion_common::Result;
     use datafusion_execution::TaskContext;
+    use datafusion_execution::runtime_env::RuntimeEnvBuilder;
     use datafusion_functions_aggregate::count::count_udaf;
     use datafusion_functions_aggregate::sum::sum_udaf;
     use datafusion_physical_expr::aggregate::AggregateExprBuilder;
@@ -135,7 +136,13 @@ mod tests {
             schema,
         )?);
 
-        let task_ctx = Arc::new(TaskContext::default());
+        // This test is for `GroupByMetrics`, which are maintained by
+        // `GroupedHashAggregateStream`. Use a finite memory pool so the partial
+        // aggregate does not take the initial-partial stream path.
+        let runtime = RuntimeEnvBuilder::new()
+            .with_memory_limit(10 * 1024 * 1024, 1.0)
+            .build_arc()?;
+        let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime));
         let _result =
             collect(Arc::clone(&aggregate_exec) as _, Arc::clone(&task_ctx)).await?;