fix(query): keep state column names unique when splitting aggregates (#9594)

* fix(query): keep state column names unique when splitting aggregates

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

* fix(query): probe allocated names when renaming colliding state columns

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

src/common/function/src/aggrs/aggr_wrapper.rs

@@ -22,6 +22,7 @@
 //! `foo_merge`'s input arg is the same as `foo_state`'s output, and its output is the same as `foo`'s input.
 //!
 
+use std::collections::HashSet;
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
@@ -183,7 +184,13 @@ impl StateMergeHelper {
                 )));
             }
         };
-        let mut lower_aggr_exprs = vec![];
+        // (unaliased state expression, output column name)
+        let mut lower_states: Vec<(Expr, String)> = vec![];
+        let mut lower_output_names: HashSet<String> = aggr
+            .group_expr
+            .iter()
+            .map(|e| e.qualified_name().1)
+            .collect();
         let mut upper_aggr_exprs = vec![];
 
         // group exprs for upper plan should refer to the output group expr as column from lower plan
@@ -218,9 +225,28 @@ impl StateMergeHelper {
                 params: aggr_func.params.clone(),
             };
             let expr = Expr::AggregateFunction(expr);
-            let lower_state_output_col_name = expr.schema_name().to_string();
 
-            lower_aggr_exprs.push(expr);
+            // The state column is named after the unaliased state expression, which can
+            // repeat even though the original output names are unique: a name-preserving
+            // rewrite such as `count(*)` -> `count(ts) AS count(*)` can make two aggregates
+            // share one state expression, and schema names hide casts, so
+            // `count(CAST(ts AS Int64))` is named `count(t.ts)` too. Identical states are
+            // computed once; a different state with a taken name gets a numeric suffix.
+            let lower_state_output_col_name =
+                if let Some((_, name)) = lower_states.iter().find(|(e, _)| *e == expr) {
+                    name.clone()
+                } else {
+                    let base = expr.schema_name().to_string();
+                    let mut name = base.clone();
+                    let mut suffix = 0;
+                    while lower_output_names.contains(&name) {
+                        suffix += 1;
+                        name = format!("{base}_{suffix}");
+                    }
+                    lower_output_names.insert(name.clone());
+                    lower_states.push((expr, name.clone()));
+                    name
+                };
 
             // then create the merge function using the physical expression of the original aggregate function
             let (name, human_display) = match aggr_expr {
@@ -270,7 +296,16 @@ impl StateMergeHelper {
         }
 
         let mut lower = aggr.clone();
-        lower.aggr_expr = lower_aggr_exprs;
+        lower.aggr_expr = lower_states
+            .into_iter()
+            .map(|(expr, name)| {
+                if expr.schema_name().to_string() == name {
+                    expr
+                } else {
+                    expr.alias(name)
+                }
+            })
+            .collect();
         let lower_plan = LogicalPlan::Aggregate(lower);
 
         // update aggregate's output schema

src/query/src/dist_plan/analyzer/test.rs

@@ -58,6 +58,7 @@ use table::table::numbers::NumbersTable;
 use table::{Table, TableRef};
 
 use super::*;
+use crate::optimizer::count_wildcard::CountWildcardToTimeIndexRule;
 
 fn collect_merge_scan_remote_dyn_filter_producer_ids(
     plan: &LogicalPlan,
@@ -1997,6 +1998,98 @@ fn expand_step_aggr_avg_limit() {
     assert_eq!(expected, result.to_string());
 }
 
+/// `count(*)` and `count(1)` are rewritten to `count(ts)` with their names preserved, so
+/// all three share one state expression and must share one state column.
+#[test]
+fn expand_step_aggr_count_wildcard_and_time_index() {
+    init_default_ut_logging();
+    let test_table = TestTable::table_with_name(0, "t".to_string());
+    let table_source = Arc::new(DefaultTableSource::new(Arc::new(
+        DfTableProviderAdapter::new(test_table),
+    )));
+    let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
+        .unwrap()
+        .aggregate(
+            Vec::<Expr>::new(),
+            vec![
+                datafusion::functions_aggregate::count::count_all(),
+                datafusion::functions_aggregate::expr_fn::count(lit(1i64)),
+                datafusion::functions_aggregate::expr_fn::count(col("ts")),
+            ],
+        )
+        .unwrap()
+        .project(vec![
+            col("count(*)").alias("a"),
+            col("count(Int64(1))"),
+            col("count(t.ts)").alias("b"),
+        ])
+        .unwrap()
+        .build()
+        .unwrap();
+
+    let config = ConfigOptions::default();
+    let plan = CountWildcardToTimeIndexRule.analyze(plan, &config).unwrap();
+    let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
+    let expected = [
+        "Projection: count(*) AS a, count(Int64(1)), count(t.ts) AS b",
+        "  Projection: count(*), count(Int64(1)), count(t.ts)",
+        "    Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(t.ts)) AS count(*), __count_merge(__count_state(t.ts)) AS count(Int64(1)), __count_merge(__count_state(t.ts)) AS count(t.ts)]]",
+        "      MergeScan [is_placeholder=false, remote_input=[",
+        "Aggregate: groupBy=[[]], aggr=[[__count_state(t.ts)]]",
+        "  TableScan: t",
+        "]]",
+    ]
+    .join("\n");
+    assert_eq!(expected, result.to_string());
+}
+
+/// After the `count(*)` rewrite, `count(ts)` and the casts of `ts` are different states with
+/// the same schema name, as schema names hide casts. The output alias `t.ts` and the second
+/// cast make sure a renamed state never takes a name that is already allocated.
+#[test]
+fn expand_step_aggr_count_wildcard_and_cast_time_index() {
+    init_default_ut_logging();
+    let test_table = TestTable::table_with_name(0, "t".to_string());
+    let table_source = Arc::new(DefaultTableSource::new(Arc::new(
+        DfTableProviderAdapter::new(test_table),
+    )));
+    let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
+        .unwrap()
+        .aggregate(
+            Vec::<Expr>::new(),
+            vec![
+                datafusion::functions_aggregate::count::count_all(),
+                datafusion::functions_aggregate::expr_fn::count(datafusion_expr::cast(
+                    col("ts"),
+                    DataType::Int64,
+                ))
+                .alias("t.ts"),
+                datafusion::functions_aggregate::expr_fn::count(datafusion_expr::cast(
+                    col("ts"),
+                    DataType::Utf8,
+                ))
+                .alias("c"),
+            ],
+        )
+        .unwrap()
+        .build()
+        .unwrap();
+
+    let config = ConfigOptions::default();
+    let plan = CountWildcardToTimeIndexRule.analyze(plan, &config).unwrap();
+    let result = DistPlannerAnalyzer {}.analyze(plan, &config).unwrap();
+    let expected = [
+        "Projection: count(*), t.ts, c",
+        "  Aggregate: groupBy=[[]], aggr=[[__count_merge(__count_state(t.ts)) AS count(*), __count_merge(__count_state(t.ts)_1) AS t.ts, __count_merge(__count_state(t.ts)_2) AS c]]",
+        "    MergeScan [is_placeholder=false, remote_input=[",
+        "Aggregate: groupBy=[[]], aggr=[[__count_state(t.ts), __count_state(CAST(t.ts AS Int64)) AS __count_state(t.ts)_1, __count_state(CAST(t.ts AS Utf8)) AS __count_state(t.ts)_2]]",
+        "  TableScan: t",
+        "]]",
+    ]
+    .join("\n");
+    assert_eq!(expected, result.to_string());
+}
+
 /// notice how `Limit` can still get expanded
 #[test]
 fn expand_part_col_aggr_limit() {

tests/cases/standalone/common/aggregate/multi_regions.result

@@ -131,6 +131,32 @@ select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g;
 | y | 2        | 6.0        |
 +---+----------+------------+
 
+-- count(*) and count(1) are rewritten to count(ts), sharing the same state column
+select count(*), count(1), count(ts) from t;
+
++----------+-----------------+-------------+
+| count(*) | count(Int64(1)) | count(t.ts) |
++----------+-----------------+-------------+
+| 4        | 4               | 4           |
++----------+-----------------+-------------+
+
+select count(*) as a, count(ts) as b from t;
+
++---+---+
+| a | b |
++---+---+
+| 4 | 4 |
++---+---+
+
+-- the cast is hidden from the state column name
+select count(*), count(ts::bigint) from t;
+
++----------+-------------+
+| count(*) | count(t.ts) |
++----------+-------------+
+| 4        | 4           |
++----------+-------------+
+
 -- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them
 -- SQLNESS PROTOCOL POSTGRES
 select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc;

tests/cases/standalone/common/aggregate/multi_regions.sql

@@ -54,6 +54,14 @@ insert into t values
 -- group keys derived from the partition column span regions
 select substr(host, 6, 1) as g, count(*), sum(val) from t group by g order by g;
 
+-- count(*) and count(1) are rewritten to count(ts), sharing the same state column
+select count(*), count(1), count(ts) from t;
+
+select count(*) as a, count(ts) as b from t;
+
+-- the cast is hidden from the state column name
+select count(*), count(ts::bigint) from t;
+
 -- grouping sets are computed on the frontend, only the PostgreSQL dialect parses them
 -- SQLNESS PROTOCOL POSTGRES
 select host, idc, sum(val) from t group by grouping sets ((host, idc), (host)) order by host, idc;