chore(query): reduce memory expansion in transform_aggregate_expand (#18372)

Author: sundy-li

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

@@ -88,7 +88,7 @@ impl PipelineBuilder {
             grouping_ids.push(!id & mask);
         }
 
-        self.main_pipeline.add_transformer(|| {
+        self.main_pipeline.add_accumulating_transformer(|| {
             TransformExpandGroupingSets::new(group_bys.clone(), grouping_ids.clone())
         });
         Ok(())

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

@@ -19,7 +19,7 @@ use databend_common_expression::types::NumberScalar;
 use databend_common_expression::BlockEntry;
 use databend_common_expression::DataBlock;
 use databend_common_expression::Scalar;
-use databend_common_pipeline_transforms::processors::Transform;
+use databend_common_pipeline_transforms::AccumulatingTransform;
 
 pub struct TransformExpandGroupingSets {
     group_bys: Vec<usize>,
@@ -35,10 +35,10 @@ impl TransformExpandGroupingSets {
     }
 }
 
-impl Transform for TransformExpandGroupingSets {
+impl AccumulatingTransform for TransformExpandGroupingSets {
     const NAME: &'static str = "TransformExpandGroupingSets";
 
-    fn transform(&mut self, data: DataBlock) -> Result<DataBlock> {
+    fn transform(&mut self, data: DataBlock) -> Result<Vec<DataBlock>> {
         let num_rows = data.num_rows();
         let num_group_bys = self.group_bys.len();
         let mut output_blocks = Vec::with_capacity(self.grouping_ids.len());
@@ -48,45 +48,53 @@ impl Transform for TransformExpandGroupingSets {
             .map(|i| data.get_by_offset(*i).clone())
             .collect::<Vec<_>>();
 
+        let mut entries = data
+            .columns()
+            .iter()
+            .cloned()
+            .chain(dup_group_by_cols.clone())
+            .collect::<Vec<_>>();
+
+        // all group columns should be nullable
+        for i in 0..num_group_bys {
+            let entry = unsafe {
+                let offset = self.group_bys.get_unchecked(i);
+                entries.get_unchecked_mut(*offset)
+            };
+            match entry {
+                BlockEntry::Const(_, data_type, _) => {
+                    *data_type = data_type.wrap_nullable();
+                }
+                BlockEntry::Column(column) => *column = column.clone().wrap_nullable(None),
+            };
+        }
+
         for &id in &self.grouping_ids {
             // Repeat data for each grouping set.
             let grouping_id_column = BlockEntry::new_const_column(
                 DataType::Number(NumberDataType::UInt32),
                 Scalar::Number(NumberScalar::UInt32(id as u32)),
                 num_rows,
             );
-            let mut entries = data
-                .columns()
-                .iter()
-                .cloned()
-                .chain(dup_group_by_cols.clone())
-                .chain(Some(grouping_id_column))
-                .collect::<Vec<_>>();
+            // This is a copy of entries which clones the buffer of columns
+            // So it's memory efficient
+            let mut current_group_entries = entries.clone();
+            current_group_entries.push(grouping_id_column);
+
             let bits = !id;
             for i in 0..num_group_bys {
                 let entry = unsafe {
                     let offset = self.group_bys.get_unchecked(i);
-                    entries.get_unchecked_mut(*offset)
+                    current_group_entries.get_unchecked_mut(*offset)
                 };
+                // Reset the column to be nullable
                 if bits & (1 << i) == 0 {
-                    // This column should be set to NULLs.
-                    *entry = BlockEntry::new_const_column(
-                        entry.data_type().wrap_nullable(),
-                        Scalar::Null,
-                        num_rows,
-                    )
-                } else {
-                    match entry {
-                        BlockEntry::Const(_, data_type, _) => {
-                            *data_type = data_type.wrap_nullable();
-                        }
-                        BlockEntry::Column(column) => *column = column.clone().wrap_nullable(None),
-                    };
+                    *entry = BlockEntry::new_const_column(entry.data_type(), Scalar::Null, num_rows)
                 }
             }
-            output_blocks.push(DataBlock::new(entries, num_rows));
+            output_blocks.push(DataBlock::new(current_group_entries, num_rows));
         }
 
-        DataBlock::concat(&output_blocks)
+        Ok(output_blocks)
     }
 }