two aggs on top of each other

Author: Samyak2

datafusion/physical-plan/src/coalesce_partitions.rs

@@ -433,6 +433,49 @@ mod tests {
             .collect()
     }
 
+    fn high_cardinality_partitioned_aggregate(
+        input: Arc<dyn ExecutionPlan>,
+        schema: &SchemaRef,
+        output_partitions: usize,
+    ) -> Result<(Arc<dyn ExecutionPlan>, Weak<RepartitionExec>)> {
+        let groups =
+            PhysicalGroupBy::new_single(vec![(col("a", schema)?, "a".to_string())]);
+        let count = AggregateExprBuilder::new(count_udaf(), vec![col("a", schema)?])
+            .schema(Arc::clone(schema))
+            .alias("COUNT(a)")
+            .build()
+            .map(Arc::new)?;
+
+        let partial = Arc::new(AggregateExec::try_new(
+            AggregateMode::Partial,
+            groups,
+            vec![count],
+            vec![None],
+            input,
+            Arc::clone(schema),
+        )?);
+
+        let hash_exprs = partial.output_group_expr();
+        let final_grouping_set = partial.group_expr().as_final();
+        let final_aggr_expr = partial.aggr_expr().to_vec();
+        let repartition = Arc::new(RepartitionExec::try_new(
+            partial,
+            Partitioning::Hash(hash_exprs, output_partitions),
+        )?);
+        let repartition_ref = Arc::downgrade(&repartition);
+
+        let final_partitioned = Arc::new(AggregateExec::try_new(
+            AggregateMode::FinalPartitioned,
+            final_grouping_set,
+            final_aggr_expr,
+            vec![None],
+            repartition,
+            Arc::clone(schema),
+        )?);
+
+        Ok((final_partitioned, repartition_ref))
+    }
+
     #[derive(Debug, Clone)]
     struct CountingGenerator {
         schema: SchemaRef,
@@ -503,49 +546,6 @@ mod tests {
         }
     }
 
-    fn high_cardinality_partitioned_aggregate(
-        input: Arc<dyn ExecutionPlan>,
-        schema: &SchemaRef,
-        output_partitions: usize,
-    ) -> Result<(Arc<dyn ExecutionPlan>, Weak<RepartitionExec>)> {
-        let groups =
-            PhysicalGroupBy::new_single(vec![(col("a", schema)?, "a".to_string())]);
-        let count = AggregateExprBuilder::new(count_udaf(), vec![col("a", schema)?])
-            .schema(Arc::clone(schema))
-            .alias("COUNT(a)")
-            .build()
-            .map(Arc::new)?;
-
-        let partial = Arc::new(AggregateExec::try_new(
-            AggregateMode::Partial,
-            groups,
-            vec![count],
-            vec![None],
-            input,
-            Arc::clone(schema),
-        )?);
-
-        let hash_exprs = partial.output_group_expr();
-        let final_grouping_set = partial.group_expr().as_final();
-        let final_aggr_expr = partial.aggr_expr().to_vec();
-        let repartition = Arc::new(RepartitionExec::try_new(
-            partial,
-            Partitioning::Hash(hash_exprs, output_partitions),
-        )?);
-        let repartition_ref = Arc::downgrade(&repartition);
-
-        let final_partitioned = Arc::new(AggregateExec::try_new(
-            AggregateMode::FinalPartitioned,
-            final_grouping_set,
-            final_aggr_expr,
-            vec![None],
-            repartition,
-            Arc::clone(schema),
-        )?);
-
-        Ok((final_partitioned, repartition_ref))
-    }
-
     #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
     #[ignore = "temporary diagnostic reproducer for layered cancellation delay"]
     async fn cancellation_delay_coalesce_repartition() -> Result<()> {
@@ -567,9 +567,16 @@ mod tests {
         let input = Arc::new(LazyMemoryExec::try_new(Arc::clone(&schema), generators)?);
         let input_refs = Arc::downgrade(&input);
         let output_partitions = 32;
-        let (partitioned_aggregate, repartition_ref) =
+        let (partitioned_aggregate, lower_repartition_ref) =
             high_cardinality_partitioned_aggregate(input, &schema, output_partitions)?;
-        let repartition_refs = vec![repartition_ref];
+        let lower_coalesce = Arc::new(CoalescePartitionsExec::new(partitioned_aggregate));
+        let (partitioned_aggregate, upper_repartition_ref) =
+            high_cardinality_partitioned_aggregate(
+                lower_coalesce,
+                &schema,
+                output_partitions,
+            )?;
+        let repartition_refs = vec![lower_repartition_ref, upper_repartition_ref];
         let plan = Arc::new(CoalescePartitionsExec::new(partitioned_aggregate));
 
         let handle = tokio::spawn(collect(plan, task_ctx));
@@ -582,9 +589,14 @@ mod tests {
         let drop_times = wait_for_repartition_drop_times(&repartition_refs, start).await;
         let total_elapsed = start.elapsed();
 
-        for elapsed in drop_times {
+        for (idx, elapsed) in drop_times.into_iter().enumerate() {
+            let repartition = match idx {
+                0 => "lower",
+                1 => "upper",
+                _ => "unknown",
+            };
             println!(
-                "final_partitioned_hash_repartition_drop_elapsed_ms={}",
+                "{repartition}_final_partitioned_hash_repartition_drop_elapsed_ms={}",
                 elapsed.as_millis()
             );
         }