diff --git a/src/common/function/src/aggrs/aggr_wrapper.rs b/src/common/function/src/aggrs/aggr_wrapper.rs index 710c383a834..8c0c475c265 100644 --- a/src/common/function/src/aggrs/aggr_wrapper.rs +++ b/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 = 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 diff --git a/src/query/src/dist_plan/analyzer/test.rs b/src/query/src/dist_plan/analyzer/test.rs index 24f8a9bc43a..a689d94c282 100644 --- a/src/query/src/dist_plan/analyzer/test.rs +++ b/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::::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::::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() { diff --git a/tests/cases/standalone/common/aggregate/multi_regions.result b/tests/cases/standalone/common/aggregate/multi_regions.result index 79ae85fbc04..bdfc5ed2dd1 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.result +++ b/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; diff --git a/tests/cases/standalone/common/aggregate/multi_regions.sql b/tests/cases/standalone/common/aggregate/multi_regions.sql index 4f99e1ed465..9c8d8af4490 100644 --- a/tests/cases/standalone/common/aggregate/multi_regions.sql +++ b/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;