feat(planner): support common table expression(CTE)

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

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
+
 use common_ast::ast::TableAlias;
 use common_ast::parser::token::Token;
 use common_ast::DisplayError;
@@ -25,6 +27,7 @@ use common_exception::Result;
 use super::AggregateInfo;
 use crate::sql::common::IndexType;
 use crate::sql::normalize_identifier;
+use crate::sql::optimizer::SExpr;
 use crate::sql::plans::Scalar;
 use crate::sql::NameResolutionContext;
 
@@ -54,7 +57,7 @@ pub enum NameResolutionResult {
 }
 
 /// `BindContext` stores all the free variables in a query and tracks the context of binding procedure.
-#[derive(Clone, Default, Debug)]
+#[derive(Clone, Debug)]
 pub struct BindContext {
     pub parent: Option<Box<BindContext>>,
 
@@ -69,20 +72,37 @@ pub struct BindContext {
 
     /// Format type of query output.
     pub format: Option<String>,
+
+    pub ctes_map: HashMap<String, CteInfo>,
+}
+
+#[derive(Clone, Debug)]
+pub struct CteInfo {
+    pub columns_alias: Vec<String>,
+    pub s_expr: SExpr,
+    pub bind_context: BindContext,
 }
 
 impl BindContext {
     pub fn new() -> Self {
-        Self::default()
+        Self {
+            parent: None,
+            columns: Vec::new(),
+            aggregate_info: AggregateInfo::default(),
+            in_grouping: false,
+            format: None,
+            ctes_map: HashMap::new(),
+        }
     }
 
     pub fn with_parent(parent: Box<BindContext>) -> Self {
         BindContext {
-            parent: Some(parent),
+            parent: Some(parent.clone()),
             columns: vec![],
             aggregate_info: Default::default(),
             in_grouping: false,
             format: None,
+            ctes_map: parent.ctes_map.clone(),
         }
     }
 
@@ -267,3 +287,9 @@ impl BindContext {
         DataSchemaRefExt::create(fields)
     }
 }
+
+impl Default for BindContext {
+    fn default() -> Self {
+        BindContext::new()
+    }
+}

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

@@ -43,7 +43,7 @@ use crate::sql::BindContext;
 impl<'a> Binder {
     pub(in crate::sql::planner::binder) async fn bind_copy(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &CopyStmt<'a>,
     ) -> Result<Plan> {
         match (&stmt.src, &stmt.dst) {
@@ -313,7 +313,7 @@ impl<'a> Binder {
     #[allow(clippy::too_many_arguments)]
     async fn bind_copy_from_table_into_stage(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &CopyStmt<'a>,
         src_catalog_name: &str,
         src_database_name: &str,
@@ -359,7 +359,7 @@ impl<'a> Binder {
     #[allow(clippy::too_many_arguments)]
     async fn bind_copy_from_table_into_uri(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &CopyStmt<'a>,
         src_catalog_name: &str,
         src_database_name: &str,
@@ -409,7 +409,7 @@ impl<'a> Binder {
     /// Bind COPY INFO <stage_location> FROM <query>
     async fn bind_copy_from_query_into_stage(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &CopyStmt<'a>,
         src_query: &Query<'_>,
         dst_stage: &str,
@@ -439,7 +439,7 @@ impl<'a> Binder {
     #[allow(clippy::too_many_arguments)]
     async fn bind_copy_from_query_into_uri(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &CopyStmt<'a>,
         src_query: &Query<'_>,
         dst_uri_location: &UriLocation,

query/src/sql/planner/binder/ddl/database.rs

@@ -48,7 +48,7 @@ use crate::sql::BindContext;
 impl<'a> Binder {
     pub(in crate::sql::planner::binder) async fn bind_show_databases(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &ShowDatabasesStmt<'a>,
     ) -> Result<Plan> {
         let ShowDatabasesStmt { limit } = stmt;

query/src/sql/planner/binder/ddl/table.rs

@@ -118,7 +118,7 @@ impl SelectBuilder {
 impl<'a> Binder {
     pub(in crate::sql::planner::binder) async fn bind_show_tables(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &ShowTablesStmt<'a>,
     ) -> Result<Plan> {
         let ShowTablesStmt {
@@ -257,7 +257,7 @@ impl<'a> Binder {
 
     pub(in crate::sql::planner::binder) async fn bind_show_tables_status(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &ShowTablesStatusStmt<'a>,
     ) -> Result<Plan> {
         let ShowTablesStatusStmt {
@@ -357,8 +357,9 @@ impl<'a> Binder {
             }
             (None, Some(query)) => {
                 // `CREATE TABLE AS SELECT ...` without column definitions
-                let init_bind_context = BindContext::new();
-                let (_s_expr, bind_context) = self.bind_query(&init_bind_context, query).await?;
+                let mut init_bind_context = BindContext::new();
+                let (_s_expr, bind_context) =
+                    self.bind_query(&mut init_bind_context, query).await?;
                 let fields = bind_context
                     .columns
                     .iter()
@@ -375,8 +376,9 @@ impl<'a> Binder {
                 // e.g. `CREATE TABLE t (i INT) AS SELECT * from old_t` with columns speicified
                 let (source_schema, source_coments) =
                     self.analyze_create_table_schema(source).await?;
-                let init_bind_context = BindContext::new();
-                let (_s_expr, bind_context) = self.bind_query(&init_bind_context, query).await?;
+                let mut init_bind_context = BindContext::new();
+                let (_s_expr, bind_context) =
+                    self.bind_query(&mut init_bind_context, query).await?;
                 let query_fields: Vec<DataField> = bind_context
                     .columns
                     .iter()
@@ -443,9 +445,9 @@ impl<'a> Binder {
             table_meta,
             cluster_keys,
             as_select: if let Some(query) = as_query {
-                let bind_context = BindContext::new();
+                let mut bind_context = BindContext::new();
                 let stmt = Statement::Query(Box::new(*query.clone()));
-                let select_plan = self.bind_statement(&bind_context, &stmt).await?;
+                let select_plan = self.bind_statement(&mut bind_context, &stmt).await?;
                 // Don't enable distributed optimization for `CREATE TABLE ... AS SELECT ...` for now
                 let opt_ctx = Arc::new(OptimizerContext::new(OptimizerConfig::default()));
                 let optimized_plan = optimize(self.ctx.clone(), opt_ctx, select_plan)?;

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

@@ -49,7 +49,7 @@ impl QueryASTIRVisitor<HashSet<String>> for DeleteCollectPushDowns {
 impl<'a> Binder {
     pub(in crate::sql::planner::binder) async fn bind_delete(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         table_reference: &'a TableReference<'a>,
         selection: &'a Option<Expr<'a>>,
     ) -> Result<Plan> {

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

@@ -62,7 +62,7 @@ use crate::sql::MetadataRef;
 impl<'a> Binder {
     pub(in crate::sql::planner::binder) async fn bind_insert(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &InsertStmt<'a>,
     ) -> Result<Plan> {
         let InsertStmt {

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

@@ -45,7 +45,7 @@ impl<'a> Binder {
     #[async_recursion]
     pub(super) async fn bind_join(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         join: &Join<'a>,
     ) -> Result<(SExpr, BindContext)> {
         let (left_child, left_context) =

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

@@ -96,14 +96,14 @@ impl<'a> Binder {
     }
 
     pub async fn bind(mut self, stmt: &Statement<'a>) -> Result<Plan> {
-        let init_bind_context = BindContext::new();
-        self.bind_statement(&init_bind_context, stmt).await
+        let mut init_bind_context = BindContext::new();
+        self.bind_statement(&mut init_bind_context, stmt).await
     }
 
     #[async_recursion::async_recursion]
     async fn bind_statement(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &Statement<'a>,
     ) -> Result<Plan> {
         let plan = match stmt {
@@ -317,7 +317,7 @@ impl<'a> Binder {
 
     async fn bind_rewrite_to_query(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         query: &str,
         rewrite_kind_r: RewriteKind,
     ) -> Result<Plan> {

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

@@ -31,6 +31,7 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 
 use crate::sql::binder::scalar_common::split_conjunctions;
+use crate::sql::binder::CteInfo;
 use crate::sql::optimizer::SExpr;
 use crate::sql::planner::binder::scalar::ScalarBinder;
 use crate::sql::planner::binder::BindContext;
@@ -56,7 +57,7 @@ pub struct SelectItem<'a> {
 impl<'a> Binder {
     pub(super) async fn bind_select_stmt(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         stmt: &SelectStmt<'a>,
         order_by: &[OrderByExpr<'a>],
     ) -> Result<(SExpr, BindContext)> {
@@ -149,14 +150,15 @@ impl<'a> Binder {
         let mut output_context = BindContext::new();
         output_context.parent = from_context.parent;
         output_context.columns = from_context.columns;
+        output_context.ctes_map = from_context.ctes_map;
 
         Ok((s_expr, output_context))
     }
 
     #[async_recursion]
     pub(crate) async fn bind_set_expr(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         set_expr: &SetExpr,
         order_by: &[OrderByExpr],
     ) -> Result<(SExpr, BindContext)> {
@@ -176,11 +178,29 @@ impl<'a> Binder {
         }
     }
 
+    #[async_recursion]
     pub(crate) async fn bind_query(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         query: &Query<'_>,
     ) -> Result<(SExpr, BindContext)> {
+        if let Some(with) = &query.with {
+            for cte in with.ctes.iter() {
+                let table_name = cte.alias.name.name.clone();
+                if bind_context.ctes_map.contains_key(&table_name) {
+                    return Err(ErrorCode::SemanticError(format!(
+                        "duplicate cte {table_name}"
+                    )));
+                }
+                let (s_expr, cte_bind_context) = self.bind_query(bind_context, &cte.query).await?;
+                let cte_info = CteInfo {
+                    columns_alias: cte.alias.columns.iter().map(|c| c.name.clone()).collect(),
+                    s_expr,
+                    bind_context: cte_bind_context.clone(),
+                };
+                bind_context.ctes_map.insert(table_name, cte_info);
+            }
+        }
         let (mut s_expr, mut bind_context) = match query.body {
             SetExpr::Select(_) | SetExpr::Query(_) => {
                 self.bind_set_expr(bind_context, &query.body, &query.order_by)
@@ -250,7 +270,7 @@ impl<'a> Binder {
 
     pub(super) async fn bind_set_operator(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         left: &SetExpr<'_>,
         right: &SetExpr<'_>,
         op: &SetOperator,

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

@@ -23,7 +23,7 @@ use crate::sql::Binder;
 impl<'a> Binder {
     pub(in crate::sql::planner::binder) async fn bind_show_functions(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         limit: &Option<ShowLimit<'a>>,
     ) -> Result<Plan> {
         // rewrite show functions to select * from system.functions ...
@@ -51,7 +51,7 @@ impl<'a> Binder {
 
     pub(in crate::sql::planner::binder) async fn bind_show_settings(
         &mut self,
-        bind_context: &BindContext,
+        bind_context: &mut BindContext,
         like: &Option<String>,
     ) -> Result<Plan> {
         let sub_query = like