feat(planner): logical plan for rcte

src/frontend/planner_test/tests/testdata/input/recursive_cte.yaml

@@ -1,12 +1,16 @@
 - name: basic
   sql: WITH RECURSIVE t1 AS (SELECT 1 AS a UNION ALL SELECT a + 1 FROM t1 WHERE a < 10) SELECT * FROM t1;
   expected_outputs:
-  - planner_error
+  - logical_plan
 - name: output column follows lhs
   sql: WITH RECURSIVE t1 AS (SELECT 1 AS a UNION ALL SELECT a + 1 FROM t1 WHERE a < 10) SELECT a FROM t1;
   expected_outputs:
-  - planner_error
+  - logical_plan
+- name: with normal column
+  sql: WITH RECURSIVE t(a) AS (VALUES(1) UNION ALL SELECT a + 1 FROM t WHERE a < 100) SELECT * FROM t;
+  expected_outputs:
+  - logical_plan
 - name: name a is leaked outside
   sql: WITH RECURSIVE t1 AS (SELECT 1 AS a UNION ALL SELECT a + 1 FROM t1 WHERE a < 10) SELECT a;
   expected_outputs:
-  - binder_error
+  - binder_error

src/frontend/planner_test/tests/testdata/output/recursive_cte.yaml

@@ -1,14 +1,33 @@
 # This file is automatically generated. See `src/frontend/planner_test/README.md` for more information.
 - name: basic
   sql: WITH RECURSIVE t1 AS (SELECT 1 AS a UNION ALL SELECT a + 1 FROM t1 WHERE a < 10) SELECT * FROM t1;
-  planner_error: |-
-    Feature is not yet implemented: recursive CTE is not supported
-    Tracking issue: https://github.com/risingwavelabs/risingwave/issues/15135
+  logical_plan: |-
+    LogicalProject { exprs: [$expr1] }
+    └─LogicalRecursiveUnion { id: 4 }
+      ├─LogicalProject { exprs: [1:Int32] }
+      │ └─LogicalValues { rows: [[]], schema: Schema { fields: [] } }
+      └─LogicalProject { exprs: [(1:Int32 + 1:Int32) as $expr1] }
+        └─LogicalFilter { predicate: (1:Int32 < 10:Int32) }
+          └─LogicalCteRef { share_id: 0 }
 - name: output column follows lhs
   sql: WITH RECURSIVE t1 AS (SELECT 1 AS a UNION ALL SELECT a + 1 FROM t1 WHERE a < 10) SELECT a FROM t1;
-  planner_error: |-
-    Feature is not yet implemented: recursive CTE is not supported
-    Tracking issue: https://github.com/risingwavelabs/risingwave/issues/15135
+  logical_plan: |-
+    LogicalProject { exprs: [$expr1] }
+    └─LogicalRecursiveUnion { id: 4 }
+      ├─LogicalProject { exprs: [1:Int32] }
+      │ └─LogicalValues { rows: [[]], schema: Schema { fields: [] } }
+      └─LogicalProject { exprs: [(1:Int32 + 1:Int32) as $expr1] }
+        └─LogicalFilter { predicate: (1:Int32 < 10:Int32) }
+          └─LogicalCteRef { share_id: 0 }
+- name: with normal column
+  sql: WITH RECURSIVE t(a) AS (VALUES(1) UNION ALL SELECT a + 1 FROM t WHERE a < 100) SELECT * FROM t;
+  logical_plan: |-
+    LogicalProject { exprs: [$expr1] }
+    └─LogicalRecursiveUnion { id: 3 }
+      ├─LogicalValues { rows: [[1:Int32]], schema: Schema { fields: [*VALUES*_0.column_0:Int32] } }
+      └─LogicalProject { exprs: [(*VALUES*_0.column_0 + 1:Int32) as $expr1] }
+        └─LogicalFilter { predicate: (*VALUES*_0.column_0 < 100:Int32) }
+          └─LogicalCteRef { share_id: 0 }
 - name: name a is leaked outside
   sql: WITH RECURSIVE t1 AS (SELECT 1 AS a UNION ALL SELECT a + 1 FROM t1 WHERE a < 10) SELECT a;
   binder_error: |

src/frontend/src/binder/bind_context.rs

@@ -94,7 +94,7 @@ pub enum BindingCteState {
     #[default]
     Init,
     /// We know the schema form after the base term resolved.
-    BaseResolved { schema: Schema },
+    BaseResolved { base: BoundSetExpr },
     /// We get the whole bound result of the (recursive) CTE.
     Bound {
         query: Either<BoundQuery, RecursiveUnion>,

src/frontend/src/binder/mod.rs

@@ -49,8 +49,9 @@ pub use insert::BoundInsert;
 use pgwire::pg_server::{Session, SessionId};
 pub use query::BoundQuery;
 pub use relation::{
-    BoundBaseTable, BoundJoin, BoundShare, BoundSource, BoundSystemTable, BoundWatermark,
-    BoundWindowTableFunction, Relation, ResolveQualifiedNameError, WindowTableFunctionKind,
+    BoundBackCteRef, BoundBaseTable, BoundJoin, BoundShare, BoundSource, BoundSystemTable,
+    BoundWatermark, BoundWindowTableFunction, Relation, ResolveQualifiedNameError,
+    WindowTableFunctionKind,
 };
 pub use select::{BoundDistinct, BoundSelect};
 pub use set_expr::*;

src/frontend/src/binder/query.rs

@@ -423,9 +423,7 @@ impl Binder {
         // reference: <https://www.postgresql.org/docs/16/sql-select.html#:~:text=the%20recursive%20self%2Dreference%20must%20appear%20on%20the%20right%2Dhand%20side%20of%20the%20UNION>
         let mut base = self.bind_set_expr(left)?;
 
-        entry.borrow_mut().state = BindingCteState::BaseResolved {
-            schema: base.schema().clone(),
-        };
+        entry.borrow_mut().state = BindingCteState::BaseResolved { base: base.clone() };
 
         // Reset context for right side, but keep `cte_to_relation`.
         let new_context = std::mem::take(&mut self.context);

src/frontend/src/binder/relation/cte_ref.rs

@@ -13,14 +13,14 @@
 // limitations under the License.
 
 use crate::binder::statement::RewriteExprsRecursive;
-use crate::binder::ShareId;
+use crate::binder::{BoundSetExpr, ShareId};
 
 /// A CTE reference, currently only used in the back reference of recursive CTE.
 /// For the non-recursive one, see [`BoundShare`](super::BoundShare).
 #[derive(Debug, Clone)]
 pub struct BoundBackCteRef {
-    #[expect(dead_code)]
     pub(crate) share_id: ShareId,
+    pub(crate) base: BoundSetExpr,
 }
 
 impl RewriteExprsRecursive for BoundBackCteRef {

src/frontend/src/binder/relation/mod.rs

@@ -26,7 +26,6 @@ use risingwave_sqlparser::ast::{
 use thiserror::Error;
 use thiserror_ext::AsReport;
 
-use self::cte_ref::BoundBackCteRef;
 use super::bind_context::ColumnBinding;
 use super::statement::RewriteExprsRecursive;
 use crate::binder::bind_context::{BindingCte, BindingCteState};
@@ -43,6 +42,7 @@ mod table_or_source;
 mod watermark;
 mod window_table_function;
 
+pub use cte_ref::BoundBackCteRef;
 pub use join::BoundJoin;
 pub use share::BoundShare;
 pub use subquery::BoundSubquery;
@@ -374,13 +374,13 @@ impl Binder {
                 BindingCteState::Init => {
                     Err(ErrorCode::BindError("Base term of recursive CTE not found, consider writing it to left side of the `UNION ALL` operator".to_string()).into())
                 }
-                BindingCteState::BaseResolved { schema } => {
+                BindingCteState::BaseResolved { base } => {
                     self.bind_table_to_context(
-                        schema.fields.iter().map(|f| (false, f.clone())),
+                        base.schema().fields.iter().map(|f| (false, f.clone())),
                         table_name.clone(),
                         Some(original_alias),
                     )?;
-                    Ok(Relation::BackCteRef(Box::new(BoundBackCteRef { share_id })))
+                    Ok(Relation::BackCteRef(Box::new(BoundBackCteRef { share_id, base })))
                 }
                 BindingCteState::Bound { query } => {
                     let schema = match &query {

src/frontend/src/optimizer/optimizer_context.rs

@@ -15,17 +15,20 @@
 use core::convert::Into;
 use core::fmt::Formatter;
 use std::cell::{RefCell, RefMut};
+use std::collections::HashMap;
 use std::marker::PhantomData;
 use std::rc::Rc;
 use std::sync::Arc;
 
 use risingwave_sqlparser::ast::{ExplainOptions, ExplainType};
 
+use crate::binder::ShareId;
 use crate::expr::{CorrelatedId, SessionTimezone};
 use crate::handler::HandlerArgs;
 use crate::optimizer::plan_node::PlanNodeId;
 use crate::session::SessionImpl;
 use crate::utils::{OverwriteOptions, WithOptions};
+use crate::PlanRef;
 
 const RESERVED_ID_NUM: u16 = 10000;
 
@@ -58,6 +61,9 @@ pub struct OptimizerContext {
     /// Store the configs can be overwritten in with clause
     /// if not specified, use the value from session variable.
     overwrite_options: OverwriteOptions,
+    /// Store the mapping between `share_id` and the corresponding
+    /// `PlanRef`, used by rcte's planning. (e.g., in `LogicalCteRef`)
+    rcte_cache: RefCell<HashMap<ShareId, PlanRef>>,
 
     _phantom: PhantomUnsend,
 }
@@ -91,6 +97,7 @@ impl OptimizerContext {
             next_expr_display_id: RefCell::new(RESERVED_ID_NUM.into()),
             total_rule_applied: RefCell::new(0),
             overwrite_options,
+            rcte_cache: RefCell::new(HashMap::new()),
             _phantom: Default::default(),
         }
     }
@@ -113,6 +120,7 @@ impl OptimizerContext {
             next_expr_display_id: RefCell::new(0),
             total_rule_applied: RefCell::new(0),
             overwrite_options: OverwriteOptions::default(),
+            rcte_cache: RefCell::new(HashMap::new()),
             _phantom: Default::default(),
         }
         .into()
@@ -230,6 +238,14 @@ impl OptimizerContext {
     pub fn get_session_timezone(&self) -> String {
         self.session_timezone.borrow().timezone()
     }
+
+    pub fn get_rcte_cache_plan(&self, id: &ShareId) -> Option<PlanRef> {
+        self.rcte_cache.borrow().get(id).cloned()
+    }
+
+    pub fn insert_rcte_cache_plan(&self, id: ShareId, plan: PlanRef) {
+        self.rcte_cache.borrow_mut().insert(id, plan);
+    }
 }
 
 impl std::fmt::Debug for OptimizerContext {

src/frontend/src/optimizer/plan_node/generic/cte_ref.rs

@@ -0,0 +1,90 @@
+// Copyright 2024 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::hash::Hash;
+
+use itertools::Itertools;
+use pretty_xmlish::{Pretty, StrAssocArr};
+use risingwave_common::catalog::Schema;
+
+use super::{impl_distill_unit_from_fields, GenericPlanNode, GenericPlanRef};
+use crate::binder::ShareId;
+use crate::optimizer::property::FunctionalDependencySet;
+use crate::{optimizer, OptimizerContextRef};
+
+#[derive(Clone, Debug)]
+pub struct CteRef<PlanRef> {
+    share_id: ShareId,
+    base: PlanRef,
+}
+
+impl<PlanRef> PartialEq for CteRef<PlanRef> {
+    fn eq(&self, other: &Self) -> bool {
+        self.share_id == other.share_id
+    }
+}
+
+impl<PlanRef> Eq for CteRef<PlanRef> {}
+
+impl<PlanRef> Hash for CteRef<PlanRef> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.share_id.hash(state);
+    }
+}
+
+impl<PlanRef> CteRef<PlanRef> {
+    pub fn new(share_id: ShareId, base: PlanRef) -> Self {
+        Self { share_id, base }
+    }
+}
+
+impl<PlanRef: GenericPlanRef> CteRef<PlanRef> {
+    pub fn get_cte_ref(&self) -> Option<optimizer::plan_node::PlanRef> {
+        self.base.ctx().get_rcte_cache_plan(&self.share_id)
+    }
+}
+
+impl<PlanRef: GenericPlanRef> GenericPlanNode for CteRef<PlanRef> {
+    fn schema(&self) -> Schema {
+        self.base.schema().clone()
+    }
+
+    fn stream_key(&self) -> Option<Vec<usize>> {
+        if let Some(plan_ref) = self.get_cte_ref() {
+            plan_ref
+                .stream_key()
+                .map(|s| s.iter().map(|i| i.to_owned()).collect_vec())
+        } else {
+            self.base
+                .stream_key()
+                .map(|s| s.iter().map(|i| i.to_owned()).collect_vec())
+        }
+    }
+
+    fn ctx(&self) -> OptimizerContextRef {
+        self.base.ctx()
+    }
+
+    fn functional_dependency(&self) -> FunctionalDependencySet {
+        self.base.functional_dependency().clone()
+    }
+}
+
+impl<PlanRef: GenericPlanRef> CteRef<PlanRef> {
+    pub fn fields_pretty<'a>(&self) -> StrAssocArr<'a> {
+        vec![("share_id", Pretty::debug(&self.share_id))]
+    }
+}
+
+impl_distill_unit_from_fields! {CteRef, GenericPlanRef}

src/frontend/src/optimizer/plan_node/generic/mod.rs

@@ -72,6 +72,10 @@ mod limit;
 pub use limit::*;
 mod max_one_row;
 pub use max_one_row::*;
+mod cte_ref;
+pub use cte_ref::*;
+mod recursive_union;
+pub use recursive_union::*;
 
 pub trait DistillUnit {
     fn distill_with_name<'a>(&self, name: impl Into<Cow<'a, str>>) -> XmlNode<'a>;

src/frontend/src/optimizer/plan_node/generic/recursive_union.rs

@@ -0,0 +1,62 @@
+// Copyright 2024 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use pretty_xmlish::StrAssocArr;
+use risingwave_common::catalog::Schema;
+
+use super::{impl_distill_unit_from_fields, GenericPlanNode, GenericPlanRef};
+use crate::optimizer::property::FunctionalDependencySet;
+use crate::OptimizerContextRef;
+
+/// `RecursiveUnion` returns the union of the rows of its inputs.
+/// note: if `all` is false, it needs to eliminate duplicates.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct RecursiveUnion<PlanRef> {
+    pub base: PlanRef,
+    pub recursive: PlanRef,
+}
+
+impl<PlanRef: GenericPlanRef> GenericPlanNode for RecursiveUnion<PlanRef> {
+    fn functional_dependency(&self) -> FunctionalDependencySet {
+        self.recursive.functional_dependency().clone()
+    }
+
+    fn schema(&self) -> Schema {
+        self.recursive.schema().clone()
+    }
+
+    fn stream_key(&self) -> Option<Vec<usize>> {
+        let fields_len = self.base.schema().len();
+        let base = self.base.stream_key();
+        if let Some(base) = base {
+            let mut base = base.to_vec();
+            base.push(fields_len);
+            Some(base)
+        } else {
+            None
+        }
+    }
+
+    fn ctx(&self) -> OptimizerContextRef {
+        self.recursive.ctx()
+    }
+}
+
+impl<PlanRef: GenericPlanRef> RecursiveUnion<PlanRef> {
+    pub fn fields_pretty<'a>(&self) -> StrAssocArr<'a> {
+        vec![]
+    }
+}
+
+impl_distill_unit_from_fields!(RecursiveUnion, GenericPlanRef);