planner: fixing wrong result after applying predicate push down for CTEs (#47891) (#48193)

close #47881

Author: ti-chi-bot

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",

planner/core/issuetest/BUILD.bazel

@@ -6,6 +6,6 @@ go_test(
     srcs = ["planner_issue_test.go"],
     flaky = True,
     race = "on",
-    shard_count = 6,
+    shard_count = 7,
     deps = ["//testkit"],
 )

planner/core/issuetest/planner_issue_test.go

@@ -115,3 +115,65 @@ func TestIssue46083(t *testing.T) {
 	tk.MustExec("CREATE TEMPORARY TABLE v0(v1 int)")
 	tk.MustExec("INSERT INTO v0 WITH ta2 AS (TABLE v0) TABLE ta2 FOR UPDATE OF ta2;")
 }
+
+func TestIssue47881(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t (id int,name varchar(10));")
+	tk.MustExec("insert into t values(1,'tt');")
+	tk.MustExec("create table t1(id int,name varchar(10),name1 varchar(10),name2 varchar(10));")
+	tk.MustExec("insert into t1 values(1,'tt','ttt','tttt'),(2,'dd','ddd','dddd');")
+	tk.MustExec("create table t2(id int,name varchar(10),name1 varchar(10),name2 varchar(10),`date1` date);")
+	tk.MustExec("insert into t2 values(1,'tt','ttt','tttt','2099-12-31'),(2,'dd','ddd','dddd','2099-12-31');")
+	rs := tk.MustQuery(`WITH bzzs AS (
+		SELECT 
+		  count(1) AS bzn 
+		FROM 
+		  t c
+	      ), 
+	      tmp1 AS (
+		SELECT 
+		  t1.* 
+		FROM 
+		  t1 
+		  LEFT JOIN bzzs ON 1 = 1 
+		WHERE 
+		  name IN ('tt') 
+		  AND bzn <> 1
+	      ), 
+	      tmp2 AS (
+		SELECT 
+		  tmp1.*, 
+		  date('2099-12-31') AS endate 
+		FROM 
+		  tmp1
+	      ), 
+	      tmp3 AS (
+		SELECT 
+		  * 
+		FROM 
+		  tmp2 
+		WHERE 
+		  endate > CURRENT_DATE 
+		UNION ALL 
+		SELECT 
+		  '1' AS id, 
+		  'ss' AS name, 
+		  'sss' AS name1, 
+		  'ssss' AS name2, 
+		  date('2099-12-31') AS endate 
+		FROM 
+		  bzzs t1 
+		WHERE 
+		  bzn = 1
+	      ) 
+	      SELECT 
+		c2.id, 
+		c3.id 
+	      FROM 
+		t2 db 
+		LEFT JOIN tmp3 c2 ON c2.id = '1' 
+		LEFT JOIN tmp3 c3 ON c3.id = '1';`)
+	rs.Check(testkit.Rows("1 1", "1 1"))
+}

planner/core/logical_plan_builder.go

@@ -4491,13 +4491,21 @@ func (b *PlanBuilder) tryBuildCTE(ctx context.Context, tn *ast.TableName, asName
 			}
 
 			if cte.cteClass == nil {
-				cte.cteClass = &CTEClass{IsDistinct: cte.isDistinct, seedPartLogicalPlan: cte.seedLP,
-					recursivePartLogicalPlan: cte.recurLP, IDForStorage: cte.storageID,
-					optFlag: cte.optFlag, HasLimit: hasLimit, LimitBeg: limitBeg,
-					LimitEnd: limitEnd, pushDownPredicates: make([]expression.Expression, 0), ColumnMap: make(map[string]*expression.Column)}
+				cte.cteClass = &CTEClass{
+					IsDistinct:               cte.isDistinct,
+					seedPartLogicalPlan:      cte.seedLP,
+					recursivePartLogicalPlan: cte.recurLP,
+					IDForStorage:             cte.storageID,
+					optFlag:                  cte.optFlag,
+					HasLimit:                 hasLimit,
+					LimitBeg:                 limitBeg,
+					LimitEnd:                 limitEnd,
+					pushDownPredicates:       make([]expression.Expression, 0),
+					ColumnMap:                make(map[string]*expression.Column),
+				}
 			}
 			var p LogicalPlan
-			lp := LogicalCTE{cteAsName: tn.Name, cteName: tn.Name, cte: cte.cteClass, seedStat: cte.seedStat, isOuterMostCTE: !b.buildingCTE}.Init(b.ctx, b.getSelectOffset())
+			lp := LogicalCTE{cteAsName: tn.Name, cteName: tn.Name, cte: cte.cteClass, seedStat: cte.seedStat}.Init(b.ctx, b.getSelectOffset())
 			prevSchema := cte.seedLP.Schema().Clone()
 			lp.SetSchema(getResultCTESchema(cte.seedLP.Schema(), b.ctx.GetSessionVars()))
 

planner/core/logical_plans.go

@@ -2006,6 +2006,7 @@ type CTEClass struct {
 	// pushDownPredicates may be push-downed by different references.
 	pushDownPredicates []expression.Expression
 	ColumnMap          map[string]*expression.Column
+	isOuterMostCTE     bool
 }
 
 const emptyCTEClassSize = int64(unsafe.Sizeof(CTEClass{}))
@@ -2037,11 +2038,10 @@ func (cc *CTEClass) MemoryUsage() (sum int64) {
 type LogicalCTE struct {
 	logicalSchemaProducer
 
-	cte            *CTEClass
-	cteAsName      model.CIStr
-	cteName        model.CIStr
-	seedStat       *property.StatsInfo
-	isOuterMostCTE bool
+	cte       *CTEClass
+	cteAsName model.CIStr
+	cteName   model.CIStr
+	seedStat  *property.StatsInfo
 }
 
 // LogicalCTETable is for CTE table

planner/core/optimizer.go

@@ -158,6 +158,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
 }
 

planner/core/recheck_cte.go

@@ -0,0 +1,53 @@
+// Copyright 2023 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/planner/funcdep"
+
+// 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 checks whether the CTE is referenced only by the main query or not.
+func RecheckCTE(p LogicalPlan) {
+	visited := funcdep.NewFastIntSet()
+	findCTEs(p, &visited, true)
+}
+
+func findCTEs(
+	p LogicalPlan,
+	visited *funcdep.FastIntSet,
+	isRootTree bool,
+) {
+	if cteReader, ok := p.(*LogicalCTE); ok {
+		cte := cteReader.cte
+		if !isRootTree {
+			// Set it to false since it's referenced by other CTEs.
+			cte.isOuterMostCTE = false
+		}
+		if visited.Has(cte.IDForStorage) {
+			return
+		}
+		visited.Insert(cte.IDForStorage)
+		// Set it when we meet it first time.
+		cte.isOuterMostCTE = isRootTree
+		findCTEs(cte.seedPartLogicalPlan, visited, false)
+		if cte.recursivePartLogicalPlan != nil {
+			findCTEs(cte.recursivePartLogicalPlan, visited, false)
+		}
+		return
+	}
+	for _, child := range p.Children() {
+		findCTEs(child, visited, isRootTree)
+	}
+}

planner/core/rule_predicate_push_down.go

@@ -985,7 +985,7 @@ func (p *LogicalCTE) PredicatePushDown(predicates []expression.Expression, _ *lo
 		// Doesn't support recursive CTE yet.
 		return predicates, p.self
 	}
-	if !p.isOuterMostCTE {
+	if !p.cte.isOuterMostCTE {
 		return predicates, p.self
 	}
 	pushedPredicates := make([]expression.Expression, len(predicates))

planner/optimize.go

@@ -504,6 +504,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)