From b584d52e26b929960f179c118bacecb7f4b02bc3 Mon Sep 17 00:00:00 2001
From: Yiding Cui <winoros@gmail.com>
Date: Wed, 1 Nov 2023 01:25:10 +0800
Subject: [PATCH] fix the value of isOuterMostCTE

---
 pkg/planner/core/BUILD.bazel             |  1 +
 pkg/planner/core/logical_plan_builder.go |  4 --
 pkg/planner/core/optimizer.go            |  3 ++
 pkg/planner/core/recheck_cte.go          | 58 ++++++++++++++++++++++++
 pkg/planner/optimize.go                  |  2 +
 5 files changed, 64 insertions(+), 4 deletions(-)
 create mode 100644 pkg/planner/core/recheck_cte.go

diff --git a/pkg/planner/core/BUILD.bazel b/pkg/planner/core/BUILD.bazel
index 8a42db25ac49f..4fdc27e94063b 100644
--- a/pkg/planner/core/BUILD.bazel
+++ b/pkg/planner/core/BUILD.bazel
@@ -44,6 +44,7 @@ go_library(
         "point_get_plan.go",
         "preprocess.go",
         "property_cols_prune.go",
+        "recheck_cte.go",
         "resolve_indices.go",
         "rule_aggregation_elimination.go",
         "rule_aggregation_push_down.go",
diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go
index 949b58044dc38..3072594750d92 100644
--- a/pkg/planner/core/logical_plan_builder.go
+++ b/pkg/planner/core/logical_plan_builder.go
@@ -4869,14 +4869,10 @@ func (b *PlanBuilder) tryBuildCTE(ctx context.Context, tn *ast.TableName, asName
 					LimitEnd:                 limitEnd,
 					pushDownPredicates:       make([]expression.Expression, 0),
 					ColumnMap:                make(map[string]*expression.Column),
-					isOuterMostCTE:           true,
 				}
 			}
 			var p LogicalPlan
 			lp := LogicalCTE{cteAsName: tn.Name, cteName: tn.Name, cte: cte.cteClass, seedStat: cte.seedStat}.Init(b.ctx, b.getSelectOffset())
-			if b.buildingCTE {
-				lp.cte.isOuterMostCTE = false
-			}
 			prevSchema := cte.seedLP.Schema().Clone()
 			lp.SetSchema(getResultCTESchema(cte.seedLP.Schema(), b.ctx.GetSessionVars()))
 
diff --git a/pkg/planner/core/optimizer.go b/pkg/planner/core/optimizer.go
index 74786afb6b166..2ab93b491ac89 100644
--- a/pkg/planner/core/optimizer.go
+++ b/pkg/planner/core/optimizer.go
@@ -180,6 +180,9 @@ func BuildLogicalPlanForTest(ctx context.Context, sctx sessionctx.Context, node
 	if err != nil {
 		return nil, nil, err
 	}
+	if logic, ok := p.(LogicalPlan); ok {
+		RecheckCTE(logic)
+	}
 	return p, p.OutputNames(), err
 }
 
diff --git a/pkg/planner/core/recheck_cte.go b/pkg/planner/core/recheck_cte.go
new file mode 100644
index 0000000000000..c5329a8e8d959
--- /dev/null
+++ b/pkg/planner/core/recheck_cte.go
@@ -0,0 +1,58 @@
+// Copyright 2016 PingCAP, Inc.
+//
+// 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.
+
+package core
+
+import "github.com/pingcap/tidb/pkg/util/intset"
+
+// RecheckCTE fills the IsOuterMostCTE field for CTEs.
+// It's a temp solution to before we fully use the Sequence to optimize the CTEs.
+// This func will find the dependent relation of each CTEs.
+func RecheckCTE(p LogicalPlan) {
+	checked := intset.NewFastIntSet()
+	ctes := make(map[int]*CTEClass)
+	inDegreeMap := make(map[int]int)
+	findCTEs(p, &checked, ctes, true, inDegreeMap)
+	for id, cte := range ctes {
+		cte.isOuterMostCTE = inDegreeMap[id] == 0
+	}
+}
+
+func findCTEs(
+	p LogicalPlan,
+	checked *intset.FastIntSet,
+	ctes map[int]*CTEClass,
+	isRootTree bool,
+	inDegree map[int]int,
+) {
+	if cteReader, ok := p.(*LogicalCTE); ok {
+		cte := cteReader.cte
+		if !isRootTree {
+			inDegree[cte.IDForStorage]++
+		}
+		if checked.Has(cte.IDForStorage) {
+			return
+		}
+		ctes[cte.IDForStorage] = cte
+		checked.Insert(cte.IDForStorage)
+		findCTEs(cte.seedPartLogicalPlan, checked, ctes, false, inDegree)
+		if cte.recursivePartLogicalPlan != nil {
+			findCTEs(cte.recursivePartLogicalPlan, checked, ctes, false, inDegree)
+		}
+		return
+	}
+	for _, child := range p.Children() {
+		findCTEs(child, checked, ctes, isRootTree, inDegree)
+	}
+}
diff --git a/pkg/planner/optimize.go b/pkg/planner/optimize.go
index 5e2f3454f163a..cb85ff4ab9536 100644
--- a/pkg/planner/optimize.go
+++ b/pkg/planner/optimize.go
@@ -520,6 +520,8 @@ func optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 		return p, names, 0, nil
 	}
 
+	core.RecheckCTE(logic)
+
 	// Handle the logical plan statement, use cascades planner if enabled.
 	if sessVars.GetEnableCascadesPlanner() {
 		finalPlan, cost, err := cascades.DefaultOptimizer.FindBestPlan(sctx, logic)