executor: fix a bug that cte.iterOutTbl did not close correctly (#26129…

…) (#26148)
pingcap · Jul 19, 2021 · d8d4578 · d8d4578
1 parent 6cae8ae
commit d8d4578
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 14 deletions.
diff --git a/executor/builder.go b/executor/builder.go
@@ -4265,17 +4265,23 @@ func (b *executorBuilder) buildCTE(v *plannercore.PhysicalCTE) Executor {
 		return nil
 	}
 
-	// 2. Build iterInTbl.
+	// 2. Build tables to store intermediate results.
 	chkSize := b.ctx.GetSessionVars().MaxChunkSize
 	tps := seedExec.base().retFieldTypes
-	iterOutTbl := cteutil.NewStorageRowContainer(tps, chkSize)
-	if err := iterOutTbl.OpenAndRef(); err != nil {
-		b.err = err
-		return nil
-	}
-
 	var resTbl cteutil.Storage
 	var iterInTbl cteutil.Storage
+	var iterOutTbl cteutil.Storage
+
+	if v.RecurPlan != nil {
+		// For non-recursive CTE, the result will be put into resTbl directly.
+		// So no need to build iterOutTbl.
+		iterOutTbl := cteutil.NewStorageRowContainer(tps, chkSize)
+		if err := iterOutTbl.OpenAndRef(); err != nil {
+			b.err = err
+			return nil
+		}
+	}
+
 	storageMap, ok := b.ctx.GetSessionVars().StmtCtx.CTEStorageMap.(map[int]*CTEStorages)
 	if !ok {
 		b.err = errors.New("type assertion for CTEStorageMap failed")

diff --git a/executor/cte.go b/executor/cte.go
@@ -142,13 +142,18 @@ func (e *CTEExec) Next(ctx context.Context, req *chunk.Chunk) (err error) {
 	if !e.resTbl.Done() {
 		resAction := setupCTEStorageTracker(e.resTbl, e.ctx, e.memTracker, e.diskTracker)
 		iterInAction := setupCTEStorageTracker(e.iterInTbl, e.ctx, e.memTracker, e.diskTracker)
-		iterOutAction := setupCTEStorageTracker(e.iterOutTbl, e.ctx, e.memTracker, e.diskTracker)
+		var iterOutAction *chunk.SpillDiskAction
+		if e.iterOutTbl != nil {
+			iterOutAction = setupCTEStorageTracker(e.iterOutTbl, e.ctx, e.memTracker, e.diskTracker)
+		}
 
 		failpoint.Inject("testCTEStorageSpill", func(val failpoint.Value) {
 			if val.(bool) && config.GetGlobalConfig().OOMUseTmpStorage {
 				defer resAction.WaitForTest()
 				defer iterInAction.WaitForTest()
-				defer iterOutAction.WaitForTest()
+				if iterOutAction != nil {
+					defer iterOutAction.WaitForTest()
+				}
 			}
 		})
 
@@ -196,13 +201,13 @@ func (e *CTEExec) Close() (err error) {
 		if err = e.recursiveExec.Close(); err != nil {
 			return err
 		}
+		// `iterInTbl` and `resTbl` are shared by multiple operators,
+		// so will be closed when the SQL finishes.
+		if err = e.iterOutTbl.DerefAndClose(); err != nil {
+			return err
+		}
 	}
 
-	// `iterInTbl` and `resTbl` are shared by multiple operators,
-	// so will be closed when the SQL finishes.
-	if err = e.iterOutTbl.DerefAndClose(); err != nil {
-		return err
-	}
 	return e.baseExecutor.Close()
 }
 

diff --git a/executor/cte_test.go b/executor/cte_test.go
@@ -115,6 +115,9 @@ func (test *CTETestSuite) TestBasicCTE(c *check.C) {
 	rows.Check(testkit.Rows("1"))
 	rows = tk.MustQuery("SELECT * FROM t1 dt WHERE EXISTS( WITH RECURSIVE qn AS (SELECT a*0 AS b UNION ALL SELECT b+1 FROM qn WHERE b=0 or b = 1) SELECT * FROM qn WHERE b=a );")
 	rows.Check(testkit.Rows("1", "2"))
+
+	rows = tk.MustQuery("with recursive  c(p) as (select 1), cte(a, b) as (select 1, 1 union select a+1, 1 from cte, c where a < 5)  select * from cte order by 1, 2;")
+	rows.Check(testkit.Rows("1 1", "2 1", "3 1", "4 1", "5 1"))
 }
 
 func (test *CTESerialTestSuite) TestSpillToDisk(c *check.C) {