fix

Author: forsaken628

src/query/catalog/src/plan/agg_index.rs

@@ -51,7 +51,7 @@ pub struct AggIndexInfo {
 }
 
 /// This meta just indicate the block is from aggregating index.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub struct AggIndexMeta {
     pub is_agg: bool,
     // Number of aggregation functions.
@@ -75,6 +75,6 @@ local_block_meta_serde!(AggIndexMeta);
 #[typetag::serde(name = "agg_index_meta")]
 impl BlockMetaInfo for AggIndexMeta {
     fn clone_self(&self) -> Box<dyn BlockMetaInfo> {
-        Box::new(self.clone())
+        Box::new(*self)
     }
 }

src/query/expression/src/aggregate/aggregate_function_state.rs

@@ -228,6 +228,10 @@ pub struct StatesLayout {
 }
 
 impl StatesLayout {
+    pub fn num_aggr_func(&self) -> usize {
+        self.states_loc.len()
+    }
+
     pub fn serialize_builders(&self, num_rows: usize) -> Vec<ColumnBuilder> {
         self.serialize_type
             .iter()

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

@@ -150,7 +150,7 @@ impl TransformPartialAggregate {
                             .params
                             .states_layout
                             .as_ref()
-                            .map(|layout| layout.states_loc.len())
+                            .map(|layout| layout.num_aggr_func())
                             .unwrap_or(0);
                         (
                             vec![],

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

@@ -93,18 +93,19 @@ impl AccumulatingTransform for PartialSingleStateAggregator {
             self.first_block_start = Some(Instant::now());
         }
 
-        let is_agg_index_block = block
+        let meta = block
             .get_meta()
             .and_then(AggIndexMeta::downcast_ref_from)
-            .map(|index| index.is_agg)
-            .unwrap_or_default();
+            .copied();
 
         let block = block.consume_convert_to_full();
-        if is_agg_index_block {
+        if let Some(meta) = meta
+            && meta.is_agg
+        {
+            assert_eq!(self.states_layout.num_aggr_func(), meta.num_agg_funcs);
             // Aggregation states are in the back of the block.
-            let states_indices = (block.num_columns() - self.states_layout.states_loc.len()
-                ..block.num_columns())
-                .collect::<Vec<_>>();
+            let start = block.num_columns() - self.states_layout.num_aggr_func();
+            let states_indices = (start..block.num_columns()).collect::<Vec<_>>();
             let states = ProjectedBlock::project(&states_indices, &block);
 
             for ((place, func), state) in self

src/query/sql/src/planner/optimizer/optimizers/operator/aggregate/stats_aggregate.rs

@@ -129,51 +129,56 @@ impl RuleStatsAggregateOptimizer {
                     for (need_rewrite_agg, agg) in
                         need_rewrite_aggs.iter().zip(agg.aggregate_functions.iter())
                     {
-                        if matches!(agg.scalar, ScalarExpr::UDAFCall(_)) {
-                            agg_results.push(agg.clone());
-                            continue;
-                        }
-
-                        let agg_func = AggregateFunction::try_from(agg.scalar.clone())?;
-
-                        if let Some((col_id, name)) = need_rewrite_agg {
-                            if let Some(stat) = stats.get(col_id) {
-                                let value_bound = if name.eq_ignore_ascii_case("min") {
-                                    &stat.min
-                                } else {
-                                    &stat.max
-                                };
-                                if !value_bound.may_be_truncated {
-                                    let scalar = ScalarExpr::ConstantExpr(ConstantExpr {
-                                        span: agg.scalar.span(),
-                                        value: value_bound.value.clone(),
-                                    });
-
-                                    let scalar =
-                                        scalar.unify_to_data_type(agg_func.return_type.as_ref());
-
-                                    eval_scalar_results.push(ScalarItem {
-                                        index: agg.index,
-                                        scalar,
-                                    });
-                                    continue;
+                        let column = if let ScalarExpr::UDAFCall(udaf) = &agg.scalar {
+                            ColumnBindingBuilder::new(
+                                udaf.display_name.clone(),
+                                agg.index,
+                                udaf.return_type.clone(),
+                                Visibility::Visible,
+                            )
+                            .build()
+                        } else {
+                            let agg_func = AggregateFunction::try_from(agg.scalar.clone())?;
+                            if let Some((col_id, name)) = need_rewrite_agg {
+                                if let Some(stat) = stats.get(col_id) {
+                                    let value_bound = if name.eq_ignore_ascii_case("min") {
+                                        &stat.min
+                                    } else {
+                                        &stat.max
+                                    };
+                                    if !value_bound.may_be_truncated {
+                                        let scalar = ScalarExpr::ConstantExpr(ConstantExpr {
+                                            span: agg.scalar.span(),
+                                            value: value_bound.value.clone(),
+                                        });
+
+                                        let scalar = scalar
+                                            .unify_to_data_type(agg_func.return_type.as_ref());
+
+                                        eval_scalar_results.push(ScalarItem {
+                                            index: agg.index,
+                                            scalar,
+                                        });
+                                        continue;
+                                    }
                                 }
                             }
-                        }
+                            ColumnBindingBuilder::new(
+                                agg_func.display_name.clone(),
+                                agg.index,
+                                agg_func.return_type.clone(),
+                                Visibility::Visible,
+                            )
+                            .build()
+                        };
 
                         // Add other aggregate functions as derived column,
                         // this will be used in aggregating index rewrite.
                         eval_scalar_results.push(ScalarItem {
                             index: agg.index,
                             scalar: ScalarExpr::BoundColumnRef(BoundColumnRef {
                                 span: agg.scalar.span(),
-                                column: ColumnBindingBuilder::new(
-                                    agg_func.display_name.clone(),
-                                    agg.index,
-                                    agg_func.return_type.clone(),
-                                    Visibility::Visible,
-                                )
-                                .build(),
+                                column,
                             }),
                         });
                         agg_results.push(agg.clone());

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

@@ -112,5 +112,7 @@ SELECT MAX(a), MIN(b), weighted_avg(a,b) from t group by b;
 2 2 1.3333334
 1 1 1.0
 
-# fix me
-# SELECT MAX(a), MIN(b), weighted_avg(a,b) from t;
+query IIR
+SELECT MAX(a), MIN(b), weighted_avg(a,b) from t;
+----
+2 1 1.2857143