refactor: refactor cte binder and fix materialized cte used in subquery (#16785)

* refactor: refactor cte binder

* extract cte context and refactor join parts, todo: set op, merge into and r cte

* all cte tests pass

* cargo fix

* make tpcds q95 materialized cte work

* remove used count in cte info

* fix set_cte_context

* fix the distinct number for ctes

* try to fix stackoverflow

* fix with consume

* add test for tpcds95

Author: xudong963

src/query/service/src/interpreters/interpreter_replace.rs

@@ -214,8 +214,6 @@ impl ReplaceInterpreter {
                 &name_resolution_ctx,
                 metadata,
                 &[],
-                Default::default(),
-                Default::default(),
             );
             let (scalar, _) = scalar_binder.bind(expr)?;
             let columns = scalar.used_columns();

src/query/sql/src/planner/binder/aggregate.rs

@@ -639,8 +639,6 @@ impl Binder {
                 &self.name_resolution_ctx,
                 self.metadata.clone(),
                 &[],
-                self.m_cte_bound_ctx.clone(),
-                self.ctes_map.clone(),
             );
             let (scalar_expr, _) = scalar_binder
                 .bind(expr)

src/query/sql/src/planner/binder/bind_context.rs

@@ -14,7 +14,9 @@
 
 use std::collections::btree_map;
 use std::collections::BTreeMap;
+use std::collections::HashMap;
 use std::hash::Hash;
+use std::sync::Arc;
 
 use dashmap::DashMap;
 use databend_common_ast::ast::Identifier;
@@ -40,6 +42,9 @@ use crate::binder::column_binding::ColumnBinding;
 use crate::binder::window::WindowInfo;
 use crate::binder::ColumnBindingBuilder;
 use crate::normalize_identifier;
+use crate::optimizer::SExpr;
+use crate::plans::MaterializedCte;
+use crate::plans::RelOperator;
 use crate::plans::ScalarExpr;
 use crate::plans::ScalarItem;
 use crate::ColumnSet;
@@ -117,10 +122,7 @@ pub struct BindContext {
 
     pub windows: WindowInfo,
 
-    /// If the `BindContext` is created from a CTE, record the cte name
-    pub cte_name: Option<String>,
-
-    pub cte_map_ref: Box<IndexMap<String, CteInfo>>,
+    pub cte_context: CteContext,
 
     /// True if there is aggregation in current context, which means
     /// non-grouping columns cannot be referenced outside aggregation
@@ -153,15 +155,90 @@ pub struct BindContext {
     pub window_definitions: DashMap<String, WindowSpec>,
 }
 
+#[derive(Clone, Debug, Default)]
+pub struct CteContext {
+    /// If the `BindContext` is created from a CTE, record the cte name
+    pub cte_name: Option<String>,
+    /// Use `IndexMap` because need to keep the insertion order
+    /// Then wrap materialized ctes to main plan.
+    pub cte_map: Box<IndexMap<String, CteInfo>>,
+    /// Record the bound s_expr of materialized cte
+    pub m_cte_bound_s_expr: HashMap<IndexType, SExpr>,
+    pub m_cte_materialized_indexes: HashMap<IndexType, IndexType>,
+}
+
+impl CteContext {
+    pub fn set_m_cte_bound_s_expr(&mut self, cte_idx: IndexType, s_expr: SExpr) {
+        self.m_cte_bound_s_expr.insert(cte_idx, s_expr);
+    }
+
+    pub fn set_m_cte_materialized_indexes(&mut self, cte_idx: IndexType, index: IndexType) {
+        self.m_cte_materialized_indexes.insert(cte_idx, index);
+    }
+
+    // Check if the materialized cte has been bound.
+    pub fn has_bound(&self, cte_idx: IndexType) -> bool {
+        self.m_cte_bound_s_expr.contains_key(&cte_idx)
+    }
+
+    // Merge two `CteContext` into one.
+    pub fn merge(&mut self, other: CteContext) {
+        let mut merged_cte_map = IndexMap::new();
+        for (left_key, left_value) in self.cte_map.iter() {
+            if let Some(right_value) = other.cte_map.get(left_key) {
+                let mut merged_value = left_value.clone();
+                if left_value.columns.is_empty() {
+                    merged_value.columns = right_value.columns.clone()
+                }
+                merged_cte_map.insert(left_key.clone(), merged_value);
+            }
+        }
+        self.cte_map = Box::new(merged_cte_map);
+        self.m_cte_bound_s_expr.extend(other.m_cte_bound_s_expr);
+        self.m_cte_materialized_indexes
+            .extend(other.m_cte_materialized_indexes);
+    }
+
+    // Wrap materialized cte to main plan.
+    // It will be called at the end of binding.
+    pub fn wrap_m_cte(&self, mut s_expr: SExpr) -> SExpr {
+        for (_, cte_info) in self.cte_map.iter().rev() {
+            if !cte_info.materialized {
+                continue;
+            }
+            if let Some(cte_s_expr) = self.m_cte_bound_s_expr.get(&cte_info.cte_idx) {
+                let materialized_output_columns = cte_info.columns.clone();
+                s_expr = SExpr::create_binary(
+                    Arc::new(RelOperator::MaterializedCte(MaterializedCte {
+                        cte_idx: cte_info.cte_idx,
+                        materialized_output_columns,
+                        materialized_indexes: self.m_cte_materialized_indexes.clone(),
+                    })),
+                    Arc::new(s_expr),
+                    Arc::new(cte_s_expr.clone()),
+                );
+            }
+        }
+        s_expr
+    }
+
+    // Set cte context to current `BindContext`.
+    // To make sure the last `BindContext` of the whole binding phase contains `cte context`
+    // Then we can wrap materialized cte to main plan.
+    pub fn set_cte_context(&mut self, cte_context: CteContext) {
+        self.cte_map = cte_context.cte_map;
+        self.m_cte_bound_s_expr = cte_context.m_cte_bound_s_expr;
+        self.m_cte_materialized_indexes = cte_context.m_cte_materialized_indexes;
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct CteInfo {
     pub columns_alias: Vec<String>,
     pub query: Query,
     pub materialized: bool,
     pub recursive: bool,
     pub cte_idx: IndexType,
-    // Record how many times this cte is used
-    pub used_count: usize,
     // If cte is materialized, save its columns
     pub columns: Vec<ColumnBinding>,
 }
@@ -174,8 +251,7 @@ impl BindContext {
             bound_internal_columns: BTreeMap::new(),
             aggregate_info: AggregateInfo::default(),
             windows: WindowInfo::default(),
-            cte_name: None,
-            cte_map_ref: Box::default(),
+            cte_context: CteContext::default(),
             in_grouping: false,
             view_info: None,
             srfs: Vec::new(),
@@ -196,8 +272,7 @@ impl BindContext {
             bound_internal_columns: BTreeMap::new(),
             aggregate_info: Default::default(),
             windows: Default::default(),
-            cte_name: parent.cte_name,
-            cte_map_ref: parent.cte_map_ref.clone(),
+            cte_context: parent.cte_context.clone(),
             in_grouping: false,
             view_info: None,
             srfs: Vec::new(),
@@ -215,8 +290,7 @@ impl BindContext {
     pub fn replace(&self) -> Self {
         let mut bind_context = BindContext::new();
         bind_context.parent = self.parent.clone();
-        bind_context.cte_name = self.cte_name.clone();
-        bind_context.cte_map_ref = self.cte_map_ref.clone();
+        bind_context.cte_context = self.cte_context.clone();
         bind_context
     }
 

src/query/sql/src/planner/binder/bind_mutation/bind.rs

@@ -32,7 +32,6 @@ use databend_common_expression::Scalar;
 use databend_common_expression::TableSchema;
 use databend_common_expression::TableSchemaRef;
 use databend_common_expression::ROW_VERSION_COL_NAME;
-use indexmap::IndexMap;
 
 use crate::binder::bind_mutation::mutation_expression::MutationExpression;
 use crate::binder::bind_mutation::mutation_expression::MutationExpressionBindResult;
@@ -258,8 +257,6 @@ impl Binder {
             &name_resolution_ctx,
             self.metadata.clone(),
             &[],
-            HashMap::new(),
-            Box::new(IndexMap::new()),
         );
 
         // Bind matched clause columns and add update fields and exprs

src/query/sql/src/planner/binder/bind_mutation/mutation_expression.rs

@@ -39,7 +39,6 @@ use crate::optimizer::SExpr;
 use crate::optimizer::SubqueryRewriter;
 use crate::plans::BoundColumnRef;
 use crate::plans::Filter;
-use crate::plans::MaterializedCte;
 use crate::plans::MutationSource;
 use crate::plans::RelOperator;
 use crate::plans::SubqueryExpr;
@@ -121,7 +120,7 @@ impl MutationExpression {
                 }
 
                 // Wrap `LogicalMaterializedCte` to `source_expr`.
-                source_s_expr = binder.wrap_cte(source_s_expr);
+                source_s_expr = source_context.cte_context.wrap_m_cte(source_s_expr);
 
                 // When there is "update *" or "insert *", prepare all source columns.
                 let all_source_columns = Self::all_source_columns(
@@ -443,26 +442,6 @@ impl Binder {
         Ok(row_id_index)
     }
 
-    fn wrap_cte(&mut self, mut s_expr: SExpr) -> SExpr {
-        for (_, cte_info) in self.ctes_map.iter().rev() {
-            if !cte_info.materialized || cte_info.used_count == 0 {
-                continue;
-            }
-            let cte_s_expr = self.m_cte_bound_s_expr.get(&cte_info.cte_idx).unwrap();
-            let materialized_output_columns = cte_info.columns.clone();
-            s_expr = SExpr::create_binary(
-                Arc::new(RelOperator::MaterializedCte(MaterializedCte {
-                    cte_idx: cte_info.cte_idx,
-                    materialized_output_columns,
-                    materialized_indexes: self.m_cte_materialized_indexes.clone(),
-                })),
-                Arc::new(s_expr),
-                Arc::new(cte_s_expr.clone()),
-            );
-        }
-        s_expr
-    }
-
     // Recursively flatten the AND expressions.
     pub fn flatten_and_scalar_expr(scalar: &ScalarExpr) -> Vec<ScalarExpr> {
         if let ScalarExpr::FunctionCall(func) = scalar
@@ -489,8 +468,6 @@ impl Binder {
                 &self.name_resolution_ctx,
                 self.metadata.clone(),
                 &[],
-                self.m_cte_bound_ctx.clone(),
-                self.ctes_map.clone(),
             );
             let (scalar, _) = scalar_binder.bind(expr)?;
             if !self.check_allowed_scalar_expr_with_subquery(&scalar)? {

src/query/sql/src/planner/binder/bind_query/bind.rs

@@ -70,7 +70,7 @@ impl Binder {
 
         for (idx, cte) in with.ctes.iter().enumerate() {
             let table_name = self.normalize_identifier(&cte.alias.name).name;
-            if bind_context.cte_map_ref.contains_key(&table_name) {
+            if bind_context.cte_context.cte_map.contains_key(&table_name) {
                 return Err(ErrorCode::SemanticError(format!(
                     "Duplicate common table expression: {table_name}"
                 )));
@@ -87,11 +87,12 @@ impl Binder {
                 materialized: cte.materialized,
                 recursive: with.recursive,
                 cte_idx: idx,
-                used_count: 0,
                 columns: vec![],
             };
-            self.ctes_map.insert(table_name.clone(), cte_info.clone());
-            bind_context.cte_map_ref.insert(table_name, cte_info);
+            bind_context
+                .cte_context
+                .cte_map
+                .insert(table_name, cte_info);
         }
 
         Ok(())
@@ -117,8 +118,6 @@ impl Binder {
             &self.name_resolution_ctx,
             self.metadata.clone(),
             &[],
-            self.m_cte_bound_ctx.clone(),
-            self.ctes_map.clone(),
         );
         let mut order_by_items = Vec::with_capacity(query.order_by.len());
 

src/query/sql/src/planner/binder/bind_query/bind_select.rs

@@ -247,6 +247,9 @@ impl Binder {
 
         let mut output_context = BindContext::new();
         output_context.parent = from_context.parent;
+        output_context
+            .cte_context
+            .set_cte_context(from_context.cte_context.clone());
         output_context.columns = from_context.columns;
 
         Ok((s_expr, output_context))

src/query/sql/src/planner/binder/bind_query/bind_value.rs

@@ -30,7 +30,6 @@ use databend_common_expression::DataSchema;
 use databend_common_expression::DataSchemaRefExt;
 use databend_common_expression::Evaluator;
 use databend_common_functions::BUILTIN_FUNCTIONS;
-use indexmap::IndexMap;
 
 use crate::binder::wrap_cast;
 use crate::optimizer::ColumnSet;
@@ -389,8 +388,6 @@ pub fn bind_values(
         name_resolution_ctx,
         metadata.clone(),
         &[],
-        HashMap::new(),
-        Box::new(IndexMap::new()),
     );
 
     let num_values = values.len();