diff --git a/pkg/executor/cte.go b/pkg/executor/cte.go index 9928bc652b639..317f2aba79163 100644 --- a/pkg/executor/cte.go +++ b/pkg/executor/cte.go @@ -545,6 +545,16 @@ func (p *cteProducer) computeChunkHash(chk *chunk.Chunk) (sel []int, err error) hashBitMap[val] = true } } else { + // Length of p.sel is init as MaxChunkSize, but the row num of chunk may still exceeds MaxChunkSize. + // So needs to handle here to make sure len(p.sel) == chk.NumRows(). + if len(p.sel) < numRows { + tmpSel := make([]int, numRows-len(p.sel)) + for i := 0; i < len(tmpSel); i++ { + tmpSel[i] = i + len(p.sel) + } + p.sel = append(p.sel, tmpSel...) + } + // All rows is selected, sel will be [0....numRows). // e.sel is setup when building executor. sel = p.sel diff --git a/pkg/executor/cte_test.go b/pkg/executor/cte_test.go index 4133e4401ed97..0f9c5c2d07a4b 100644 --- a/pkg/executor/cte_test.go +++ b/pkg/executor/cte_test.go @@ -505,3 +505,19 @@ func TestCTEShareCorColumn(t *testing.T) { tk.MustExec("insert into t1 values(1), (2);") tk.MustQuery("SELECT * FROM t1 dt WHERE EXISTS( WITH RECURSIVE qn AS (SELECT a AS b UNION ALL SELECT b+1 FROM qn WHERE b=0 or b = 1) SELECT * FROM qn dtqn1 where exists (select /*+ NO_DECORRELATE() */ b from qn where dtqn1.b+1));").Check(testkit.Rows("1", "2")) } + +func TestCTESmallChunkSize(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test;") + tk.MustExec("drop table if exists t1") + tk.MustExec("create table t1(c1 int);") + insertStr := "insert into t1 values (0)" + for i := 1; i < 300; i++ { + insertStr += fmt.Sprintf(", (%d)", i) + } + tk.MustExec(insertStr) + tk.MustExec("set @@tidb_max_chunk_size = 32;") + tk.MustQuery("with recursive cte1(c1) as (select c1 from t1 union select c1 + 1 c1 from cte1 limit 1 offset 100) select * from cte1;").Check(testkit.Rows("100")) + tk.MustExec("set @@tidb_max_chunk_size = default;") +}