chore(query): refactor new group_agg and test

src/query/expression/src/aggregate/mod.rs

@@ -91,6 +91,10 @@ impl HashTableConfig {
     pub fn with_partial(mut self, partial_agg: bool, active_threads: usize) -> Self {
         self.partial_agg = partial_agg;
 
+        if active_threads == 0 {
+            return self;
+        }
+
         // init max_partial_capacity
         let total_shared_cache_size = active_threads * L3_CACHE_SIZE;
         let cache_per_active_thread =

src/query/service/src/pipelines/builders/builder_aggregate.rs

@@ -106,6 +106,9 @@ impl PipelineBuilder {
             .settings
             .get_enable_experimental_aggregate_hashtable()?
             && self.ctx.get_cluster().is_empty();
+        // let enable_experimental_aggregate_hashtable = self
+        //     .settings
+        //     .get_enable_experimental_aggregate_hashtable()?;
 
         let params = Self::build_aggregator_params(
             aggregate.input.output_schema()?,
@@ -215,6 +218,9 @@ impl PipelineBuilder {
             .settings
             .get_enable_experimental_aggregate_hashtable()?
             && self.ctx.get_cluster().is_empty();
+        // let enable_experimental_aggregate_hashtable = self
+        //     .settings
+        //     .get_enable_experimental_aggregate_hashtable()?;
 
         let params = Self::build_aggregator_params(
             aggregate.before_group_by_schema.clone(),

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

@@ -19,8 +19,12 @@ use bumpalo::Bump;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_expression::AggregateHashTable;
 use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::DataBlock;
+use databend_common_expression::HashTableConfig;
+use databend_common_expression::PartitionedPayload;
+use databend_common_expression::PayloadFlushState;
 use databend_common_hashtable::FastHash;
 use databend_common_hashtable::HashtableEntryMutRefLike;
 use databend_common_hashtable::HashtableEntryRefLike;
@@ -77,7 +81,8 @@ impl<Method: HashMethodBounds, V: Send + Sync + 'static> ExchangeSorting
                         AggregateMeta::Partitioned { .. } => unreachable!(),
                         AggregateMeta::Serialized(v) => Ok(v.bucket),
                         AggregateMeta::HashTable(v) => Ok(v.bucket),
-                        AggregateMeta::AggregateHashTable(_) => unreachable!(),
+                        AggregateMeta::AggregateHashTable(_) => Ok(-1),
+                        AggregateMeta::AggregatePayload(_) => todo!("AGG_HASHTABLE"),
                         AggregateMeta::Spilled(_)
                         | AggregateMeta::Spilling(_)
                         | AggregateMeta::BucketSpilled(_) => Ok(-1),
@@ -139,6 +144,34 @@ fn scatter<Method: HashMethodBounds, V: Copy + Send + Sync + 'static>(
     Ok(res)
 }
 
+// TODO: buckets and partitions have a relationship of integer division
+#[allow(dead_code)]
+fn agg_hashtable_scatter(
+    payload: PartitionedPayload,
+    buckets: usize,
+) -> Result<Vec<AggregateHashTable>> {
+    let mut buckets = Vec::with_capacity(buckets);
+
+    for _ in 0..buckets.capacity() {
+        let config = HashTableConfig::default().with_partial(true, 0);
+        buckets.push(AggregateHashTable::new(
+            payload.group_types.clone(),
+            payload.aggrs.clone(),
+            config,
+        ));
+    }
+
+    let mut state = PayloadFlushState::default();
+    let new_payload = payload.repartition(buckets.len(), &mut state);
+
+    for (idx, item) in new_payload.payloads.iter().enumerate() {
+        let mut flush_state = PayloadFlushState::default();
+        buckets[idx].combine_payload(item, &mut flush_state)?;
+    }
+
+    Ok(buckets)
+}
+
 impl<Method: HashMethodBounds, V: Copy + Send + Sync + 'static> FlightScatter
     for HashTableHashScatter<Method, V>
 {
@@ -176,8 +209,16 @@ impl<Method: HashMethodBounds, V: Copy + Send + Sync + 'static> FlightScatter
                             });
                         }
                     }
-
+                    // AggregateMeta::AggregateHashTable(payload) => {
+                    //     for agg_hashtable in agg_hashtable_scatter(payload, self.buckets)? {
+                    //         blocks.push(
+                    //         DataBlock::empty_with_meta(
+                    //                     AggregateMeta::<Method, V>::create_agg_hashtable(agg_hashtable.payload)
+                    //                 ))
+                    //     }
+                    // },
                     AggregateMeta::AggregateHashTable(_) => todo!("AGG_HASHTABLE"),
+                    AggregateMeta::AggregatePayload(_) => todo!("AGG_HASHTABLE"),
                 };
 
                 return Ok(blocks);

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

@@ -21,6 +21,7 @@ use databend_common_expression::BlockMetaInfoPtr;
 use databend_common_expression::Column;
 use databend_common_expression::DataBlock;
 use databend_common_expression::PartitionedPayload;
+use databend_common_expression::Payload;
 
 use crate::pipelines::processors::transforms::aggregator::HashTableCell;
 use crate::pipelines::processors::transforms::group_by::HashMethodBounds;
@@ -31,6 +32,12 @@ pub struct HashTablePayload<T: HashMethodBounds, V: Send + Sync + 'static> {
     pub cell: HashTableCell<T, V>,
 }
 
+pub struct AggregatePayload {
+    pub bucket: isize,
+    pub payload: Payload,
+    pub max_partition_count: usize,
+}
+
 pub struct SerializedPayload {
     pub bucket: isize,
     pub data_block: DataBlock,
@@ -54,6 +61,7 @@ pub enum AggregateMeta<Method: HashMethodBounds, V: Send + Sync + 'static> {
     Serialized(SerializedPayload),
     HashTable(HashTablePayload<Method, V>),
     AggregateHashTable(PartitionedPayload),
+    AggregatePayload(AggregatePayload),
     BucketSpilled(BucketSpilledPayload),
     Spilled(Vec<BucketSpilledPayload>),
     Spilling(HashTablePayload<PartitionedHashMethod<Method>, V>),
@@ -69,6 +77,20 @@ impl<Method: HashMethodBounds, V: Send + Sync + 'static> AggregateMeta<Method, V
         }))
     }
 
+    pub fn create_agg_payload(
+        bucket: isize,
+        payload: Payload,
+        max_partition_count: usize,
+    ) -> BlockMetaInfoPtr {
+        Box::new(AggregateMeta::<Method, V>::AggregatePayload(
+            AggregatePayload {
+                bucket,
+                payload,
+                max_partition_count,
+            },
+        ))
+    }
+
     pub fn create_agg_hashtable(payload: PartitionedPayload) -> BlockMetaInfoPtr {
         Box::new(AggregateMeta::<Method, V>::AggregateHashTable(payload))
     }
@@ -136,6 +158,9 @@ impl<Method: HashMethodBounds, V: Send + Sync + 'static> Debug for AggregateMeta
             AggregateMeta::AggregateHashTable(_) => {
                 f.debug_struct("AggregateMeta:AggHashTable").finish()
             }
+            AggregateMeta::AggregatePayload(_) => {
+                f.debug_struct("AggregateMeta:AggregatePayload").finish()
+            }
         }
     }
 }

src/query/service/src/pipelines/processors/transforms/aggregator/serde/transform_aggregate_serializer.rs

@@ -146,6 +146,7 @@ impl<Method: HashMethodBounds> TransformAggregateSerializer<Method> {
                         return Ok(Event::Sync);
                     }
                     AggregateMeta::AggregateHashTable(_) => todo!("AGG_HASHTABLE"),
+                    AggregateMeta::AggregatePayload(_) => todo!("AGG_HASHTABLE"),
                 }
             }
         }

src/query/service/src/pipelines/processors/transforms/aggregator/serde/transform_exchange_aggregate_serializer.rs

@@ -177,6 +177,7 @@ impl<Method: HashMethodBounds> BlockMetaTransform<ExchangeShuffleMeta>
                 }
 
                 Some(AggregateMeta::AggregateHashTable(_)) => todo!("AGG_HASHTABLE"),
+                Some(AggregateMeta::AggregatePayload(_)) => todo!("AGG_HASHTABLE"),
             };
         }
 

src/query/service/src/pipelines/processors/transforms/aggregator/serde/transform_exchange_group_by_serializer.rs

@@ -54,6 +54,7 @@ use futures_util::future::BoxFuture;
 use log::info;
 use opendal::Operator;
 
+use super::SerializePayload;
 use crate::api::serialize_block;
 use crate::api::ExchangeShuffleMeta;
 use crate::pipelines::processors::transforms::aggregator::exchange_defines;
@@ -207,7 +208,27 @@ impl<Method: HashMethodBounds> BlockMetaTransform<ExchangeShuffleMeta>
                         },
                     ));
                 }
-                Some(AggregateMeta::AggregateHashTable(_)) => todo!("AGG_HASHTABLE"),
+                Some(AggregateMeta::AggregateHashTable(payload)) => {
+                    if index == self.local_pos {
+                        serialized_blocks.push(FlightSerialized::DataBlock(block.add_meta(
+                            Some(Box::new(AggregateMeta::<Method, ()>::AggregateHashTable(
+                                payload,
+                            ))),
+                        )?));
+                        continue;
+                    }
+                    let bucket = -1;
+                    let mut stream = SerializeGroupByStream::create(
+                        &self.method,
+                        SerializePayload::PartitionedPayload(payload),
+                    );
+                    serialized_blocks.push(FlightSerialized::DataBlock(match stream.next() {
+                        None => DataBlock::empty(),
+                        Some(data_block) => {
+                            serialize_block(bucket, data_block?, &self.ipc_fields, &self.options)?
+                        }
+                    }));
+                }
                 Some(AggregateMeta::HashTable(payload)) => {
                     if index == self.local_pos {
                         serialized_blocks.push(FlightSerialized::DataBlock(block.add_meta(
@@ -216,15 +237,19 @@ impl<Method: HashMethodBounds> BlockMetaTransform<ExchangeShuffleMeta>
                         continue;
                     }
 
-                    let mut stream = SerializeGroupByStream::create(&self.method, payload);
-                    let bucket = stream.payload.bucket;
+                    let bucket = payload.bucket;
+                    let mut stream = SerializeGroupByStream::create(
+                        &self.method,
+                        SerializePayload::HashTablePayload(payload),
+                    );
                     serialized_blocks.push(FlightSerialized::DataBlock(match stream.next() {
                         None => DataBlock::empty(),
                         Some(data_block) => {
                             serialize_block(bucket, data_block?, &self.ipc_fields, &self.options)?
                         }
                     }));
                 }
+                Some(AggregateMeta::AggregatePayload(_)) => todo!("AGG_HASHTABLE"),
             };
         }