feat: support common table expressions (#379)

* feat: support non-recursive CTEs

* feat: support recursive CTEs

* test: cover recursive CTE operators and errors

* refactor: decouple recursive CTEs from spill

Author: KKould

src/binder/expr.rs

@@ -134,25 +134,7 @@ impl<'a, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'a, '_, T
             &mut PlanArena<'arena>,
         ) -> Result<LogicalPlan, DatabaseError>,
     {
-        let BinderContext {
-            table_cache,
-            view_cache,
-            transaction,
-            scala_functions,
-            table_functions,
-            ..
-        } = &self.context;
-        let mut binder = Binder::new(
-            BinderContext::new(
-                table_cache,
-                view_cache,
-                *transaction,
-                scala_functions,
-                table_functions,
-            ),
-            self.args,
-            Some(&self.context),
-        );
+        let mut binder = Binder::new(self.context.fork_empty(), self.args, Some(&self.context));
         let sub_query = build(&mut binder, arena)?;
         let correlated = binder.context.has_outer_refs();
         Ok((sub_query, correlated))

src/binder/mod.rs

@@ -22,6 +22,7 @@ macro_rules! with_query_bind_step {
     }};
 }
 
+#[cfg(feature = "orm")]
 pub(crate) use with_query_bind_step;
 
 pub mod aggregate;
@@ -65,7 +66,7 @@ use crate::errors::DatabaseError;
 use crate::expression::ScalarExpression;
 use crate::planner::operator::join::JoinType;
 use crate::planner::operator::mark_apply::MarkApplyQuantifier;
-use crate::planner::{LogicalPlan, PlanArena};
+use crate::planner::{LogicalPlan, PlanArena, PlanRef};
 use crate::storage::{TableCache, Transaction, ViewCache};
 use crate::types::tuple::Schema;
 use crate::types::value::DataValue;
@@ -134,6 +135,19 @@ pub struct BoundSource<'a> {
     pub(crate) source: Source<'a>,
 }
 
+#[derive(Debug, Clone)]
+pub(crate) struct CteBinding {
+    pub(crate) table_name: TableName,
+    pub(crate) depth: usize,
+    pub(crate) plan_ref: PlanRef,
+}
+
+#[derive(Clone, Copy)]
+pub(crate) struct CteCheckpoint {
+    len: usize,
+    depth: usize,
+}
+
 impl BoundSource<'_> {
     pub(crate) fn matches_name(&self, table_name: &str) -> bool {
         self.table_name.as_ref() == table_name
@@ -231,6 +245,8 @@ pub struct BinderContext<'a, T: Transaction> {
     // Tips: retain binding order so wildcard expansion and position derivation
     // follow FROM/JOIN order directly.
     pub(crate) bind_table: Vec<BoundSource<'a>>,
+    ctes: Vec<CteBinding>,
+    pub(crate) cte_depth: usize,
     // alias
     expr_aliases: BTreeMap<(Option<String>, String), ScalarExpression>,
     table_aliases: HashMap<TableName, TableName>,
@@ -295,6 +311,8 @@ impl<'a, T: Transaction> BinderContext<'a, T> {
             view_cache,
             transaction,
             bind_table: Default::default(),
+            ctes: Default::default(),
+            cte_depth: 0,
             expr_aliases: Default::default(),
             table_aliases: Default::default(),
             group_by_exprs: vec![],
@@ -319,6 +337,8 @@ impl<'a, T: Transaction> BinderContext<'a, T> {
             view_cache: self.view_cache,
             transaction: self.transaction,
             bind_table: self.bind_table.clone(),
+            ctes: self.ctes.clone(),
+            cte_depth: self.cte_depth,
             expr_aliases: self.expr_aliases.clone(),
             table_aliases: self.table_aliases.clone(),
             group_by_exprs: self.group_by_exprs.clone(),
@@ -336,13 +356,52 @@ impl<'a, T: Transaction> BinderContext<'a, T> {
     /// This is used while binding an independent input, such as the right side
     /// of a join, before merging its newly bound sources into the parent scope.
     pub(crate) fn fork_empty(&self) -> Self {
-        BinderContext::new(
+        let mut context = BinderContext::new(
             self.table_cache,
             self.view_cache,
             self.transaction,
             self.scala_functions,
             self.table_functions,
-        )
+        );
+        context.ctes = self.ctes.clone();
+        context.cte_depth = self.cte_depth;
+        context
+    }
+
+    pub(crate) fn cte(&self, table_name: &TableName) -> Option<&CteBinding> {
+        self.ctes
+            .iter()
+            .rev()
+            .find(|cte| cte.table_name == *table_name)
+    }
+
+    pub(crate) fn add_cte(&mut self, cte: CteBinding) -> Result<(), DatabaseError> {
+        if self
+            .ctes
+            .iter()
+            .rev()
+            .take_while(|binding| binding.depth == cte.depth)
+            .any(|binding| binding.table_name == cte.table_name)
+        {
+            return Err(DatabaseError::UnsupportedStmt(format!(
+                "duplicate CTE name: {}",
+                cte.table_name
+            )));
+        }
+        self.ctes.push(cte);
+        Ok(())
+    }
+
+    pub(crate) fn cte_checkpoint(&self) -> CteCheckpoint {
+        CteCheckpoint {
+            len: self.ctes.len(),
+            depth: self.cte_depth,
+        }
+    }
+
+    pub(crate) fn restore_ctes(&mut self, checkpoint: CteCheckpoint) {
+        self.ctes.truncate(checkpoint.len);
+        self.cte_depth = checkpoint.depth;
     }
 
     pub fn step(&mut self, bind_step: QueryBindStep) {

src/binder/parser.rs

@@ -16,7 +16,7 @@ use super::select::{
     BindPlanAggregated, BindPlanComplete, BindPlanDistinct, BindPlanFiltered, BindPlanFrom,
     BindPlanProjected, BindPlanSelectList, BindPlanStart, JoinConstraintInput, TableAliasInput,
 };
-use super::{is_valid_identifier, with_query_bind_step, Binder, QueryBindStep, SetOperatorKind};
+use super::{is_valid_identifier, Binder, CteBinding, QueryBindStep, SetOperatorKind};
 #[cfg(feature = "copy")]
 use crate::binder::copy::{ExtSource, FileFormat};
 use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef, TableName};
@@ -34,6 +34,7 @@ use crate::planner::operator::alter_table::change_column::{DefaultChange, NotNul
 use crate::planner::operator::join::{JoinCondition, JoinOperator as LJoinOperator, JoinType};
 use crate::planner::operator::mark_apply::MarkApplyQuantifier;
 use crate::planner::operator::project::ProjectOperator;
+use crate::planner::operator::recursive_cte::{RecursiveCteOperator, RecursiveScanOperator};
 use crate::planner::operator::sort::SortField;
 use crate::planner::operator::Operator;
 use crate::planner::{Childrens, LogicalPlan, PlanArena};
@@ -2832,15 +2833,171 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
         Ok(select_items)
     }
 
+    fn bind_cte_plan(
+        &mut self,
+        cte: &sqlparser::ast::Cte,
+        alias: &TableAliasInput,
+        arena: &mut PlanArena,
+    ) -> Result<LogicalPlan, DatabaseError> {
+        let mut binder = Binder::new(self.context.fork(), self.args, Some(&self.context));
+        let plan = binder.bind_query(&cte.query, arena)?;
+        let source_name = arena.temp_table();
+        binder.bind_alias(plan, &alias.columns, alias.name.clone(), source_name, arena)
+    }
+
+    fn add_cte_plan(
+        &mut self,
+        table_name: TableName,
+        plan: LogicalPlan,
+        arena: &mut PlanArena,
+    ) -> Result<(), DatabaseError> {
+        let plan_ref = arena.alloc_plan(plan);
+        self.context.add_cte(CteBinding {
+            table_name,
+            depth: self.context.cte_depth,
+            plan_ref,
+        })
+    }
+
+    fn bind_recursive_cte_plan(
+        &mut self,
+        cte: &sqlparser::ast::Cte,
+        alias: &TableAliasInput,
+        arena: &mut PlanArena,
+    ) -> Result<LogicalPlan, DatabaseError> {
+        if cte.query.with.is_some()
+            || cte.query.order_by.is_some()
+            || cte.query.limit_clause.is_some()
+        {
+            return Err(DatabaseError::UnsupportedStmt(
+                "recursive CTE query clauses are not supported".to_string(),
+            ));
+        }
+
+        let SetExpr::SetOperation {
+            op: SetOperator::Union,
+            set_quantifier: SetQuantifier::All,
+            left,
+            right,
+        } = cte.query.body.as_ref()
+        else {
+            return Err(DatabaseError::UnsupportedStmt(
+                "recursive CTEs require a top-level UNION ALL".to_string(),
+            ));
+        };
+
+        let mut anchor_binder = Binder::new(self.context.fork(), self.args, Some(&self.context));
+        let anchor = anchor_binder.bind_set_expr(left, arena)?;
+        if anchor
+            .referenced_table()
+            .iter()
+            .any(|table| table == &alias.name)
+        {
+            return Err(DatabaseError::UnsupportedStmt(
+                "the recursive CTE cannot be referenced by its anchor".to_string(),
+            ));
+        }
+
+        let source_name = arena.temp_table();
+        let mut anchor = anchor_binder.bind_alias(
+            anchor,
+            &alias.columns,
+            alias.name.clone(),
+            source_name,
+            arena,
+        )?;
+        let schema = anchor.output_schema(arena).clone();
+        let scan = LogicalPlan::new(
+            Operator::RecursiveScan(RecursiveScanOperator {
+                schema_ref: schema.clone(),
+            }),
+            Childrens::None,
+        );
+
+        let cte_checkpoint = self.context.cte_checkpoint();
+        self.add_cte_plan(alias.name.clone(), scan, arena)?;
+        let recursive = {
+            let mut binder = Binder::new(self.context.fork(), self.args, Some(&self.context));
+            binder.bind_set_expr(right, arena)
+        };
+        self.context.restore_ctes(cte_checkpoint);
+        let mut recursive = recursive?;
+
+        fn recursive_scan_count(plan: &LogicalPlan) -> usize {
+            usize::from(matches!(&plan.operator, Operator::RecursiveScan(_)))
+                + plan
+                    .childrens
+                    .iter()
+                    .map(recursive_scan_count)
+                    .sum::<usize>()
+        }
+
+        if recursive_scan_count(&recursive) != 1 {
+            return Err(DatabaseError::UnsupportedStmt(
+                "the recursive term must reference its CTE exactly once".to_string(),
+            ));
+        }
+
+        let recursive_schema = recursive.output_schema(arena);
+        if schema.len() != recursive_schema.len() {
+            return Err(DatabaseError::MisMatch(
+                "the anchor column count",
+                "the recursive column count",
+            ));
+        }
+        if !schema
+            .iter()
+            .zip(recursive_schema)
+            .all(|(anchor, recursive)| {
+                arena.column(*anchor).datatype() == arena.column(*recursive).datatype()
+            })
+        {
+            return Err(DatabaseError::UnsupportedStmt(
+                "recursive CTE column types must match the anchor".to_string(),
+            ));
+        }
+
+        Ok(RecursiveCteOperator::build(schema, anchor, recursive))
+    }
+
     pub(crate) fn bind_query(
         &mut self,
         query: &Query,
         arena: &mut PlanArena,
     ) -> Result<LogicalPlan, DatabaseError> {
         let origin_step = self.context.step_now();
+        let cte_checkpoint = self.context.cte_checkpoint();
+        if let Some(with) = &query.with {
+            if with.recursive && with.cte_tables.len() != 1 {
+                return Err(DatabaseError::UnsupportedStmt(
+                    "only one recursive CTE is supported".to_string(),
+                ));
+            }
+
+            self.context.cte_depth += 1;
+            for cte in &with.cte_tables {
+                if cte.from.is_some() {
+                    return Err(DatabaseError::UnsupportedStmt(
+                        "CTE FROM clauses are not supported".to_string(),
+                    ));
+                }
+                if matches!(
+                    cte.materialized,
+                    Some(sqlparser::ast::CteAsMaterialized::Materialized)
+                ) {
+                    return Err(DatabaseError::UnsupportedStmt(
+                        "materialized CTEs are not supported".to_string(),
+                    ));
+                }
 
-        if let Some(_with) = &query.with {
-            // TODO support with clause.
+                let alias = sql_table_alias(cte.alias.clone());
+                let plan = if with.recursive {
+                    self.bind_recursive_cte_plan(cte, &alias, arena)?
+                } else {
+                    self.bind_cte_plan(cte, &alias, arena)?
+                };
+                self.add_cte_plan(alias.name, plan, arena)?;
+            }
         }
 
         let order_by_exprs = if let Some(order_by) = &query.order_by {
@@ -2883,6 +3040,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
             plan = self.bind_limit(plan, limit_clause, arena)?;
         }
 
+        self.context.restore_ctes(cte_checkpoint);
         self.context.step(origin_step);
         Ok(plan)
     }
@@ -3160,6 +3318,36 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_bind_non_recursive_ctes() -> Result<(), DatabaseError> {
+        let tables = build_t1_table()?;
+        let mut arena = PlanArena::new(&tables.table_arena);
+        let mut plan = tables.plan_with_arena(
+            "with first(a) as (select c1 from t1), \
+             second as (select a from first) select a from second",
+            &mut arena,
+        )?;
+
+        let schema = plan.output_schema(&mut arena);
+        assert_eq!(schema.len(), 1);
+        assert_eq!(arena.column(schema[0]).name(), "a");
+
+        assert_unsupported(
+            tables
+                .plan("with recursive cte as (select 1) select * from cte")
+                .unwrap_err(),
+            "recursive CTEs",
+        );
+        assert_unsupported(
+            tables
+                .plan("with cte as (select 1), cte as (select 2) select * from cte")
+                .unwrap_err(),
+            "duplicate CTE name",
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn force_nest_loop_join_marks_join_operator() -> Result<(), DatabaseError> {
         let tables = build_t1_table()?;