feat(query): support aggregate common subplans in CSE

Author: SkyFan2002

src/query/sql/src/planner/optimizer/optimizers/cse/analyze.rs

@@ -159,6 +159,7 @@ mod tests {
     use databend_common_expression::TableDataType;
     use databend_common_expression::TableField;
     use databend_common_expression::TableSchema;
+    use databend_common_expression::types::DataType;
     use databend_common_expression::types::NumberDataType;
     use databend_common_meta_app::schema::CatalogInfo;
     use databend_common_meta_app::schema::DatabaseType;
@@ -167,10 +168,19 @@ mod tests {
     use databend_common_meta_app::schema::TableMeta;
 
     use super::*;
+    use crate::ColumnBindingBuilder;
+    use crate::Symbol;
+    use crate::Visibility;
     use crate::planner::metadata::Metadata;
+    use crate::plans::Aggregate;
+    use crate::plans::AggregateFunction;
+    use crate::plans::AggregateMode;
+    use crate::plans::BoundColumnRef;
     use crate::plans::Join;
     use crate::plans::JoinType;
     use crate::plans::RelOperator;
+    use crate::plans::ScalarExpr;
+    use crate::plans::ScalarItem;
     use crate::plans::Scan;
 
     #[derive(Debug)]
@@ -240,6 +250,62 @@ mod tests {
         })))
     }
 
+    fn column_expr(metadata: &Metadata, table_index: usize) -> ScalarExpr {
+        let column = metadata.columns_by_table_index(table_index)[0].clone();
+        BoundColumnRef {
+            span: None,
+            column: ColumnBindingBuilder::new(
+                column.name(),
+                column.index(),
+                Box::new(column.data_type()),
+                Visibility::Visible,
+            )
+            .table_index(Some(table_index))
+            .build(),
+        }
+        .into()
+    }
+
+    fn max_aggregate_expr(
+        metadata: &Metadata,
+        table_index: usize,
+        output_index: Symbol,
+        with_group_by: bool,
+    ) -> SExpr {
+        let group_items = if with_group_by {
+            vec![ScalarItem {
+                scalar: column_expr(metadata, table_index),
+                index: Symbol::new(output_index.as_usize() + 1),
+            }]
+        } else {
+            vec![]
+        };
+
+        SExpr::create_unary(
+            Arc::new(RelOperator::Aggregate(Aggregate {
+                mode: AggregateMode::Initial,
+                group_items,
+                aggregate_functions: vec![ScalarItem {
+                    scalar: ScalarExpr::AggregateFunction(AggregateFunction {
+                        span: None,
+                        func_name: "max".to_string(),
+                        distinct: false,
+                        params: vec![],
+                        args: vec![column_expr(metadata, table_index)],
+                        return_type: Box::new(DataType::Number(NumberDataType::UInt64)),
+                        sort_descs: vec![],
+                        display_name: "max(a)".to_string(),
+                    }),
+                    index: output_index,
+                }],
+                from_distinct: false,
+                rank_limit: None,
+                grouping_sets: None,
+            })),
+            Arc::new(scan_expr(metadata, table_index)),
+        )
+    }
+
     fn cross_join_expr(left: SExpr, right: SExpr) -> SExpr {
         SExpr::create_binary(
             Arc::new(RelOperator::Join(Join {
@@ -313,4 +379,50 @@ mod tests {
                 .all(|cte| matches!(cte.child(0).unwrap().plan(), RelOperator::Scan(_)))
         );
     }
+
+    #[test]
+    fn test_analyze_common_subexpression_matches_identical_aggregates() {
+        let mut metadata = Metadata::default();
+        let t1 = fake_fuse_table(1, "t1");
+
+        let t1_left = add_table(&mut metadata, t1.clone());
+        let t1_right = add_table(&mut metadata, t1);
+
+        let left = max_aggregate_expr(&metadata, t1_left, Symbol::new(10), false);
+        let right = max_aggregate_expr(&metadata, t1_right, Symbol::new(11), false);
+        let root = cross_join_expr(left, right);
+
+        let (replacements, materialized_ctes) =
+            analyze_common_subexpression(&root, &mut metadata).unwrap();
+
+        assert_eq!(replacements.len(), 2);
+        assert_eq!(materialized_ctes.len(), 1);
+        assert!(matches!(
+            materialized_ctes[0].child(0).unwrap().plan(),
+            RelOperator::Aggregate(_)
+        ));
+    }
+
+    #[test]
+    fn test_analyze_common_subexpression_matches_identical_group_aggregates() {
+        let mut metadata = Metadata::default();
+        let t1 = fake_fuse_table(1, "t1");
+
+        let t1_left = add_table(&mut metadata, t1.clone());
+        let t1_right = add_table(&mut metadata, t1);
+
+        let left = max_aggregate_expr(&metadata, t1_left, Symbol::new(10), true);
+        let right = max_aggregate_expr(&metadata, t1_right, Symbol::new(12), true);
+        let root = cross_join_expr(left, right);
+
+        let (replacements, materialized_ctes) =
+            analyze_common_subexpression(&root, &mut metadata).unwrap();
+
+        assert_eq!(replacements.len(), 2);
+        assert_eq!(materialized_ctes.len(), 1);
+        assert!(matches!(
+            materialized_ctes[0].child(0).unwrap().plan(),
+            RelOperator::Aggregate(_)
+        ));
+    }
 }

src/query/sql/src/planner/optimizer/optimizers/cse/table_signature.rs

@@ -14,18 +14,32 @@
 
 use std::collections::HashMap;
 
+use crate::ColumnBindingBuilder;
 use crate::ColumnEntry;
 use crate::IndexType;
+use crate::ScalarExpr;
+use crate::Symbol;
+use crate::Visibility;
 use crate::optimizer::ir::SExpr;
 use crate::planner::metadata::Metadata;
+use crate::plans::Aggregate;
 use crate::plans::Join;
 use crate::plans::JoinType;
 use crate::plans::RelOperator;
+use crate::plans::ScalarItem;
 use crate::plans::Scan;
+use crate::plans::VisitorMut;
+use crate::plans::walk_expr_mut;
 
-#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct TableSignature {
     pub tables: Vec<IndexType>,
+    pub aggregate: Option<AggregateSignature>,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct AggregateSignature {
+    pub aggregate: Aggregate,
 }
 
 pub fn collect_table_signatures(
@@ -63,6 +77,7 @@ fn collect_table_signatures_rec(
             signature_to_exprs
                 .entry(TableSignature {
                     tables: tables.clone(),
+                    aggregate: None,
                 })
                 .or_default()
                 .push((path.clone(), expr.clone()));
@@ -81,15 +96,123 @@ fn collect_table_signatures_rec(
             signature_to_exprs
                 .entry(TableSignature {
                     tables: tables.clone(),
+                    aggregate: None,
                 })
                 .or_default()
                 .push((path.clone(), expr.clone()));
             Some(tables)
         }
+        RelOperator::EvalScalar(_) if child_tables.len() == 1 && child_tables[0].is_some() => {
+            child_tables[0].clone()
+        }
+        RelOperator::Aggregate(aggregate)
+            if child_tables.len() == 1 && child_tables[0].is_some() =>
+        {
+            let tables = child_tables[0].clone().unwrap();
+            if let Some(aggregate_signature) = aggregate_signature(aggregate, expr.child(0).ok()?) {
+                signature_to_exprs
+                    .entry(TableSignature {
+                        tables: tables.clone(),
+                        aggregate: Some(aggregate_signature),
+                    })
+                    .or_default()
+                    .push((path.clone(), expr.clone()));
+            }
+            None
+        }
         _ => None,
     }
 }
 
+fn aggregate_signature(aggregate: &Aggregate, input: &SExpr) -> Option<AggregateSignature> {
+    if aggregate.rank_limit.is_some() || aggregate.grouping_sets.is_some() {
+        return None;
+    }
+
+    let input_columns = input
+        .derive_relational_prop()
+        .ok()?
+        .output_columns
+        .iter()
+        .copied()
+        .enumerate()
+        .map(|(position, column)| (column, Symbol::new(position)))
+        .collect::<HashMap<_, _>>();
+
+    let mut aggregate = aggregate.clone();
+    aggregate.group_items = normalize_scalar_items(&aggregate.group_items, &input_columns)?;
+    aggregate.aggregate_functions =
+        normalize_scalar_items(&aggregate.aggregate_functions, &input_columns)?;
+
+    Some(AggregateSignature { aggregate })
+}
+
+fn normalize_scalar_items(
+    items: &[ScalarItem],
+    input_columns: &HashMap<Symbol, Symbol>,
+) -> Option<Vec<ScalarItem>> {
+    items
+        .iter()
+        .enumerate()
+        .map(|(position, item)| {
+            Some(ScalarItem {
+                scalar: normalize_scalar_expr(&item.scalar, input_columns)?,
+                index: Symbol::new(position),
+            })
+        })
+        .collect()
+}
+
+fn normalize_scalar_expr(
+    scalar: &ScalarExpr,
+    input_columns: &HashMap<Symbol, Symbol>,
+) -> Option<ScalarExpr> {
+    let mut scalar = scalar.clone();
+    let mut visitor = NormalizeColumnVisitor { input_columns };
+    visitor.visit(&mut scalar).ok()?;
+    Some(scalar)
+}
+
+struct NormalizeColumnVisitor<'a> {
+    input_columns: &'a HashMap<Symbol, Symbol>,
+}
+
+impl VisitorMut<'_> for NormalizeColumnVisitor<'_> {
+    fn visit(&mut self, expr: &mut ScalarExpr) -> databend_common_exception::Result<()> {
+        walk_expr_mut(self, expr)
+    }
+
+    fn visit_bound_column_ref(
+        &mut self,
+        col: &mut crate::plans::BoundColumnRef,
+    ) -> databend_common_exception::Result<()> {
+        let Some(normalized) = self.input_columns.get(&col.column.index) else {
+            return Err(databend_common_exception::ErrorCode::Internal(
+                "aggregate CSE column is not produced by input",
+            ));
+        };
+        col.column = ColumnBindingBuilder::new(
+            normalized.to_string(),
+            *normalized,
+            col.column.data_type.clone(),
+            Visibility::Visible,
+        )
+        .build();
+        Ok(())
+    }
+
+    fn visit_aggregate_function(
+        &mut self,
+        aggregate: &mut crate::plans::AggregateFunction,
+    ) -> databend_common_exception::Result<()> {
+        aggregate.display_name.clear();
+        for expr in aggregate.exprs_mut() {
+            self.visit(expr)?;
+        }
+        Ok(())
+    }
+}
+
 fn scan_signature(scan: &Scan, metadata: &Metadata) -> Option<IndexType> {
     let has_internal_column = scan.columns.iter().any(|column_index| {
         let column = metadata.column(*column_index);

tests/sqllogictests/suites/mode/standalone/explain/common_subexpression_optimizer.test

@@ -677,3 +677,68 @@ drop table time_dim;
 
 statement ok
 drop table store;
+
+statement ok
+create or replace table cse_agg_t as
+select number as a, number % 3 as b
+from numbers(6);
+
+query T nosort
+explain select l.b, l.max_a, r.max_a
+from (select b, max(a) as max_a from cse_agg_t group by b) l
+join (select b, max(a) as max_a from cse_agg_t group by b) r
+on l.b = r.b;
+----
+Sequence
+├── MaterializedCTE: cte_cse_0
+│   └── AggregateFinal
+│       ├── output columns: [max(a) (#2), cse_agg_t.b (#1)]
+│       ├── group by: [b]
+│       ├── aggregate functions: [max(a)]
+│       ├── estimated rows: 3.00
+│       └── AggregatePartial
+│           ├── group by: [b]
+│           ├── aggregate functions: [max(a)]
+│           ├── estimated rows: 3.00
+│           └── TableScan
+│               ├── table: default.default.cse_agg_t
+│               ├── scan id: 0
+│               ├── output columns: [a (#0), b (#1)]
+│               ├── read rows: 6
+│               ├── read size: < 1 KiB
+│               ├── partitions total: 1
+│               ├── partitions scanned: 1
+│               ├── pruning stats: [segments: <range pruning: 1 to 1 cost: <slt:ignore>>, blocks: <range pruning: 1 to 1 cost: <slt:ignore>>]
+│               ├── push downs: [filters: [], limit: NONE]
+│               └── estimated rows: 6.00
+└── HashJoin
+    ├── output columns: [cse_agg_t.b (#1), max(a) (#2), max(a) (#5)]
+    ├── join type: INNER
+    ├── build keys: [r.b (#4)]
+    ├── probe keys: [l.b (#1)]
+    ├── keys is null equal: [false]
+    ├── filters: []
+    ├── build join filters:
+    │   └── filter id:0, build key:r.b (#4), probe targets:[l.b (#1)@scan0], filter type:bloom,inlist,min_max
+    ├── estimated rows: 9.00
+    ├── MaterializeCTERef(Build)
+    │   ├── cte_name: cte_cse_0
+    │   ├── cte_schema: [b (#4), max(a) (#5)]
+    │   └── estimated rows: 3.00
+    └── MaterializeCTERef(Probe)
+        ├── cte_name: cte_cse_0
+        ├── cte_schema: [b (#1), max(a) (#2)]
+        └── estimated rows: 3.00
+
+query III rowsort
+select l.b, l.max_a, r.max_a
+from (select b, max(a) as max_a from cse_agg_t group by b) l
+join (select b, max(a) as max_a from cse_agg_t group by b) r
+on l.b = r.b;
+----
+0 3 3
+1 4 4
+2 5 5
+
+statement ok
+drop table cse_agg_t;