refactor: move aggregation spill logic into partial aggregate (#18999)

* refactor: create a new partial

* refactor: move aggregation spill logic into partial aggregate

* refactor: move aggregation spill logic into partial aggregate

* chore: catch up main

* chore: clean up

* enable enable_experiment_aggregate for test

* chore: make writers lazy init

* make max_aggregate_spill_level to 1

* fix: final aggregate must call spill finish in every round

* fix: final aggregate must call spill finish in every round

* fix: final aggregate must call spill finish in every round

* set enable_experiment_aggregate = 0

Author: dqhl76

src/common/base/src/runtime/perf/flamegraph_main_template.html

@@ -48,8 +48,6 @@ <h1>Query Performance Flamegraphs</h1>
 
 <script>
     function resizeIframe(iframe) {
-        // A small delay can help ensure the content has been fully rendered
-        // before calculating its height, especially for complex SVGs.
         setTimeout(() => {
             try {
                 const body = iframe.contentWindow.document.body;
@@ -66,4 +64,4 @@ <h1>Query Performance Flamegraphs</h1>
     }
 </script>
 </body>
-</html>
+</html>

src/query/service/src/physical_plans/physical_aggregate_partial.rs

@@ -38,14 +38,15 @@ use databend_common_sql::IndexType;
 use databend_common_storage::DataOperator;
 use itertools::Itertools;
 
+use crate::clusters::ClusterHelper;
 use crate::physical_plans::explain::PlanStatsInfo;
 use crate::physical_plans::format::AggregatePartialFormatter;
 use crate::physical_plans::format::PhysicalFormat;
 use crate::physical_plans::physical_plan::IPhysicalPlan;
 use crate::physical_plans::physical_plan::PhysicalPlan;
 use crate::physical_plans::physical_plan::PhysicalPlanMeta;
 use crate::pipelines::processors::transforms::aggregator::AggregateInjector;
-use crate::pipelines::processors::transforms::aggregator::NewTransformAggregateSpillWriter;
+use crate::pipelines::processors::transforms::aggregator::NewTransformPartialAggregate;
 use crate::pipelines::processors::transforms::aggregator::PartialSingleStateAggregator;
 use crate::pipelines::processors::transforms::aggregator::SharedPartitionStream;
 use crate::pipelines::processors::transforms::aggregator::TransformAggregateSpillWriter;
@@ -220,51 +221,66 @@ impl IPhysicalPlan for AggregatePartial {
             });
         }
 
-        builder.main_pipeline.add_transform(|input, output| {
-            Ok(ProcessorPtr::create(TransformPartialAggregate::try_create(
-                builder.ctx.clone(),
-                input,
-                output,
-                params.clone(),
-                partial_agg_config.clone(),
-            )?))
-        })?;
+        if params.enable_experiment_aggregate {
+            let cluster = &builder.ctx.get_cluster();
+            let streams_num = if !builder.is_exchange_parent() {
+                1
+            } else {
+                cluster.nodes.len()
+            };
+            let local_pos = cluster.ordered_index();
+            let shared_partition_streams = (0..streams_num)
+                .map(|_| {
+                    SharedPartitionStream::new(
+                        builder.main_pipeline.output_len(),
+                        max_block_rows,
+                        max_block_bytes,
+                        MAX_AGGREGATE_HASHTABLE_BUCKETS_NUM as usize,
+                    )
+                })
+                .collect::<Vec<_>>();
+
+            builder.main_pipeline.add_transform(|input, output| {
+                Ok(ProcessorPtr::create(
+                    NewTransformPartialAggregate::try_create(
+                        builder.ctx.clone(),
+                        input,
+                        output,
+                        params.clone(),
+                        partial_agg_config.clone(),
+                        shared_partition_streams.clone(),
+                        local_pos,
+                    )?,
+                ))
+            })?;
+        } else {
+            builder.main_pipeline.add_transform(|input, output| {
+                Ok(ProcessorPtr::create(TransformPartialAggregate::try_create(
+                    builder.ctx.clone(),
+                    input,
+                    output,
+                    params.clone(),
+                    partial_agg_config.clone(),
+                )?))
+            })?;
+        }
 
         // If cluster mode, spill write will be completed in exchange serialize, because we need scatter the block data first
-        if !builder.is_exchange_parent() {
+        if !builder.is_exchange_parent() && !params.enable_experiment_aggregate {
             let operator = DataOperator::instance().spill_operator();
             let location_prefix = builder.ctx.query_id_spill_prefix();
-            if params.enable_experiment_aggregate {
-                let shared_partition_stream = SharedPartitionStream::new(
-                    builder.main_pipeline.output_len(),
-                    max_block_rows,
-                    max_block_bytes,
-                    MAX_AGGREGATE_HASHTABLE_BUCKETS_NUM as usize,
-                );
-                builder.main_pipeline.add_transform(|input, output| {
-                    Ok(ProcessorPtr::create(
-                        NewTransformAggregateSpillWriter::try_create(
-                            input,
-                            output,
-                            builder.ctx.clone(),
-                            shared_partition_stream.clone(),
-                        )?,
-                    ))
-                })?;
-            } else {
-                builder.main_pipeline.add_transform(|input, output| {
-                    Ok(ProcessorPtr::create(
-                        TransformAggregateSpillWriter::try_create(
-                            builder.ctx.clone(),
-                            input,
-                            output,
-                            operator.clone(),
-                            params.clone(),
-                            location_prefix.clone(),
-                        )?,
-                    ))
-                })?;
-            }
+            builder.main_pipeline.add_transform(|input, output| {
+                Ok(ProcessorPtr::create(
+                    TransformAggregateSpillWriter::try_create(
+                        builder.ctx.clone(),
+                        input,
+                        output,
+                        operator.clone(),
+                        params.clone(),
+                        location_prefix.clone(),
+                    )?,
+                ))
+            })?;
         }
 
         builder.exchange_injector = AggregateInjector::create(builder.ctx.clone(), params.clone());

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

@@ -22,7 +22,6 @@ use databend_common_expression::DataBlock;
 use databend_common_expression::PartitionedPayload;
 use databend_common_expression::Payload;
 use databend_common_expression::PayloadFlushState;
-use databend_common_expression::MAX_AGGREGATE_HASHTABLE_BUCKETS_NUM;
 use databend_common_pipeline::core::Pipeline;
 use databend_common_pipeline::core::ProcessorPtr;
 use databend_common_settings::FlightCompression;
@@ -32,8 +31,6 @@ use crate::pipelines::processors::transforms::aggregator::aggregate_meta::Aggreg
 use crate::pipelines::processors::transforms::aggregator::serde::TransformExchangeAggregateSerializer;
 use crate::pipelines::processors::transforms::aggregator::serde::TransformExchangeAsyncBarrier;
 use crate::pipelines::processors::transforms::aggregator::AggregatorParams;
-use crate::pipelines::processors::transforms::aggregator::NewTransformAggregateSpillWriter;
-use crate::pipelines::processors::transforms::aggregator::SharedPartitionStream;
 use crate::pipelines::processors::transforms::aggregator::TransformAggregateDeserializer;
 use crate::pipelines::processors::transforms::aggregator::TransformAggregateSerializer;
 use crate::pipelines::processors::transforms::aggregator::TransformAggregateSpillWriter;
@@ -117,7 +114,7 @@ fn scatter_payload(mut payload: Payload, buckets: usize) -> Result<Vec<Payload>>
     Ok(buckets)
 }
 
-fn scatter_partitioned_payload(
+pub fn scatter_partitioned_payload(
     partitioned_payload: PartitionedPayload,
     buckets: usize,
 ) -> Result<Vec<PartitionedPayload>> {
@@ -264,25 +261,7 @@ impl ExchangeInjector for AggregateInjector {
     ) -> Result<()> {
         let params = self.aggregator_params.clone();
 
-        if self.aggregator_params.enable_experiment_aggregate {
-            let shared_partition_stream = SharedPartitionStream::new(
-                pipeline.output_len(),
-                self.aggregator_params.max_block_rows,
-                self.aggregator_params.max_block_bytes,
-                MAX_AGGREGATE_HASHTABLE_BUCKETS_NUM as usize,
-            );
-
-            pipeline.add_transform(|input, output| {
-                Ok(ProcessorPtr::create(
-                    NewTransformAggregateSpillWriter::try_create(
-                        input,
-                        output,
-                        self.ctx.clone(),
-                        shared_partition_stream.clone(),
-                    )?,
-                ))
-            })?;
-        } else {
+        if !self.aggregator_params.enable_experiment_aggregate {
             let operator = DataOperator::instance().spill_operator();
             let location_prefix = self.ctx.query_id_spill_prefix();
 
@@ -322,18 +301,6 @@ impl ExchangeInjector for AggregateInjector {
             .position(|x| x == local_id)
             .unwrap();
 
-        let mut partition_streams = vec![];
-        if self.aggregator_params.enable_experiment_aggregate {
-            for _i in 0..shuffle_params.destination_ids.len() {
-                partition_streams.push(SharedPartitionStream::new(
-                    pipeline.output_len(),
-                    self.aggregator_params.max_block_rows,
-                    self.aggregator_params.max_block_bytes,
-                    MAX_AGGREGATE_HASHTABLE_BUCKETS_NUM as usize,
-                ));
-            }
-        }
-
         pipeline.add_transform(|input, output| {
             Ok(ProcessorPtr::create(
                 TransformExchangeAggregateSerializer::try_create(
@@ -345,7 +312,6 @@ impl ExchangeInjector for AggregateInjector {
                     params.clone(),
                     compression,
                     local_pos,
-                    partition_streams.clone(),
                 )?,
             ))
         })?;

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

@@ -137,7 +137,11 @@ pub enum AggregateMeta {
     BucketSpilled(BucketSpilledPayload),
     Spilled(Vec<BucketSpilledPayload>),
 
-    Partitioned { bucket: isize, data: Vec<Self> },
+    Partitioned {
+        bucket: isize,
+        data: Vec<Self>,
+        activate_worker: Option<usize>,
+    },
 
     NewBucketSpilled(NewSpilledPayload),
     NewSpilled(Vec<NewSpilledPayload>),
@@ -180,8 +184,16 @@ impl AggregateMeta {
         Box::new(AggregateMeta::BucketSpilled(payload))
     }
 
-    pub fn create_partitioned(bucket: isize, data: Vec<Self>) -> BlockMetaInfoPtr {
-        Box::new(AggregateMeta::Partitioned { data, bucket })
+    pub fn create_partitioned(
+        bucket: isize,
+        data: Vec<Self>,
+        activate_worker: Option<usize>,
+    ) -> BlockMetaInfoPtr {
+        Box::new(AggregateMeta::Partitioned {
+            data,
+            bucket,
+            activate_worker,
+        })
     }
 
     pub fn create_new_bucket_spilled(payload: NewSpilledPayload) -> BlockMetaInfoPtr {

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

@@ -29,8 +29,7 @@ pub use aggregate_exchange_injector::AggregateInjector;
 pub use aggregate_meta::*;
 pub use aggregator_params::AggregatorParams;
 pub use build_partition_bucket::build_partition_bucket;
-pub use new_aggregate::NewTransformAggregateSpillWriter;
-pub use new_aggregate::SharedPartitionStream;
+pub use new_aggregate::*;
 pub use transform_aggregate_expand::TransformExpandGroupingSets;
 pub use transform_aggregate_final::TransformFinalAggregate;
 pub use transform_aggregate_partial::TransformPartialAggregate;

src/query/service/src/pipelines/processors/transforms/aggregator/new_aggregate/datablock_splitter.rs

@@ -30,8 +30,9 @@ pub fn split_partitioned_meta_into_datablocks(
     let base_chunk_size = total_len / outputs_len;
     let remainder = total_len % outputs_len;
 
-    let mut result = Vec::with_capacity(outputs_len);
     let mut data_iter = data.into_iter();
+    let mut chunks = Vec::with_capacity(outputs_len);
+    let mut activated_workers = 0;
 
     for index in 0..outputs_len {
         let chunk_size = if index < remainder {
@@ -41,8 +42,17 @@ pub fn split_partitioned_meta_into_datablocks(
         };
 
         let chunk: Vec<AggregateMeta> = data_iter.by_ref().take(chunk_size).collect();
+        if !chunk.is_empty() {
+            activated_workers += 1;
+        }
+        chunks.push(chunk);
+    }
+
+    let activate_worker = Some(activated_workers);
+    let mut result = Vec::with_capacity(outputs_len);
+    for chunk in chunks {
         result.push(DataBlock::empty_with_meta(
-            AggregateMeta::create_partitioned(bucket, chunk),
+            AggregateMeta::create_partitioned(bucket, chunk, activate_worker),
         ));
     }
 

src/query/service/src/pipelines/processors/transforms/aggregator/new_aggregate/mod.rs

@@ -15,14 +15,14 @@
 mod datablock_splitter;
 mod new_aggregate_spiller;
 mod new_final_aggregate_state;
-mod new_transform_aggregate_spill_writer;
+mod new_transform_aggregate_partial;
 mod new_transform_final_aggregate;
 mod transform_partition_bucket_scatter;
 
 pub use datablock_splitter::split_partitioned_meta_into_datablocks;
 pub use new_aggregate_spiller::NewAggregateSpiller;
 pub use new_aggregate_spiller::SharedPartitionStream;
 pub use new_final_aggregate_state::FinalAggregateSharedState;
-pub use new_transform_aggregate_spill_writer::NewTransformAggregateSpillWriter;
+pub use new_transform_aggregate_partial::NewTransformPartialAggregate;
 pub use new_transform_final_aggregate::NewFinalAggregateTransform;
 pub use transform_partition_bucket_scatter::TransformPartitionBucketScatter;