fix(query): try fix hang for cluster aggregate (#18655)

* fix(query): try fix hang for cluster aggregate

* fix(query): try fix hang for cluster aggregate

* fix(query): try fix hang for cluster aggregate

* fix(query): try fix hang for cluster aggregate

* fix(query): try fix hang for cluster aggregate

* fix(query): try fix hang for cluster aggregate

* fix(query): try fix hang for cluster aggregate

Author: zhang2014

src/query/service/src/physical_plans/physical_aggregate_final.rs

@@ -186,11 +186,19 @@ impl IPhysicalPlan for AggregateFinal {
 
         self.input.build_pipeline(builder)?;
 
+        // For distributed plans, since we are unaware of the data size processed by other nodes,
+        // we estimate the parallelism based on the worst-case scenario.
+        let after_group_parallel = match self.input.is_distributed_plan() {
+            true => builder.settings.get_max_threads()? as usize,
+            false => builder.main_pipeline.output_len(),
+        };
+
         builder.exchange_injector = old_inject;
         build_partition_bucket(
             &mut builder.main_pipeline,
             params.clone(),
             max_restore_worker,
+            after_group_parallel,
         )
     }
 }

src/query/service/src/pipelines/processors/transforms/aggregator/new_transform_partition_bucket.rs

@@ -580,6 +580,7 @@ pub fn build_partition_bucket(
     pipeline: &mut Pipeline,
     params: Arc<AggregatorParams>,
     max_restore_worker: u64,
+    after_worker: usize,
 ) -> Result<()> {
     let input_nums = pipeline.output_len();
     let transform = NewTransformPartitionBucket::create(input_nums, params.clone())?;
@@ -610,7 +611,7 @@ pub fn build_partition_bucket(
         )?))
     })?;
 
-    pipeline.try_resize(input_nums)?;
+    pipeline.try_resize(after_worker)?;
 
     Ok(())
 }

src/query/service/src/servers/flight/v1/exchange/exchange_transform.rs

@@ -48,7 +48,9 @@ impl ExchangeTransform {
 
                 // exchange writer sink and resize and exchange reader
                 let len = params.destination_ids.len();
-                let local_pipe = if params.allow_adjust_parallelism {
+                let local_pipe = if params.allow_adjust_parallelism
+                    && params.exchange_injector.exchange_sorting().is_none()
+                {
                     ctx.get_settings().get_max_threads()? as usize
                 } else {
                     1

tests/sqllogictests/suites/ee/02_ee_aggregating_index/02_0001_agg_index_projected_scan.test

@@ -106,11 +106,11 @@ export function finish(state) {
 }
 $$;
 
-query IIR
+query I rowsort
 SELECT MAX(a), MIN(b), weighted_avg(a,b) from t group by b;
 ----
-2 2 1.3333334
 1 1 1.0
+2 2 1.3333334
 
 query IIR
 SELECT MAX(a), MIN(b), weighted_avg(a,b) from t;

tests/sqllogictests/suites/query/cte/basic_r_cte.test

@@ -254,7 +254,7 @@ select concat('城市',rn::varchar) city from t1 where rn<=5;
 
 statement ok
 insert into train
-select concat('G',row_number()over()::varchar),c1.city,c2.city, n from city c1, city c2, (select 600 n union select 800 union select 1200 union select 1600) a ;
+select concat('G',row_number()over()::varchar),city_1,city_2, n from (SELECT c1.city as city_1,c2.city as city_2, n FROM city c1, city c2, (select 600 n union select 800 union select 1200 union select 1600) a order by c1.city,c2.city, n);
 
 statement ok
 insert into passenger
@@ -281,7 +281,7 @@ select
 from
         t0,(select 1 n union all select 2);
 ----
-236700 473200 210000
+224100 448000 210000
 
 statement ok
 CREATE OR REPLACE table t_tree (id int, parent int, name string);

tests/sqllogictests/suites/query/window_function/window_bound.test

@@ -267,23 +267,23 @@ SELECT a, DENSE_RANK() OVER (ORDER BY a ROWS BETWEEN UNBOUNDED PRECEDING AND UNB
 6 4
 7 5
 
-query I
+query I rowsort
 SELECT DISTINCT lead((861, FALSE, FALSE), 9, (849, TRUE, TRUE)) OVER (
     PARTITION BY 15560425903542832284, 965871850213131579 
     ORDER BY 13746504519650342222, 5897530378272856518 ASC NULLS FIRST) 
 FROM range(100, 12000000, 467);
 ----
-(861,0,0)
 (849,1,1)
+(861,0,0)
 
-query II
+query II rowsort
 SELECT DISTINCT lead((861, FALSE, FALSE), 9, (849, TRUE, TRUE)) OVER (
     PARTITION BY 15560425903542832284, 965871850213131579 
     ORDER BY 13746504519650342222, 5897530378272856518 ASC NULLS FIRST) 
 FROM range(100, 120000000, 467);
 ----
-(861,0,0)
 (849,1,1)
+(861,0,0)
 
 statement ok
 DROP DATABASE test_window_bound;