From 81a93a6977e58b0eba2341277037118c62c015f9 Mon Sep 17 00:00:00 2001
From: tiancaiamao <tiancaiamao@gmail.com>
Date: Wed, 24 Aug 2022 12:38:21 +0800
Subject: [PATCH] executor, util: fix UnionScan Next() skip reading data when
 passed chunk capacity is 0 (#36961)

close pingcap/tidb#36903
---
 executor/union_scan.go      |  4 ++++
 executor/union_scan_test.go | 13 +++++++++++++
 util/chunk/chunk.go         |  6 +++++-
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/executor/union_scan.go b/executor/union_scan.go
index ef1077dfa0580..9c0483d974f1c 100644
--- a/executor/union_scan.go
+++ b/executor/union_scan.go
@@ -137,7 +137,11 @@ func (us *UnionScanExec) open(ctx context.Context) error {
 func (us *UnionScanExec) Next(ctx context.Context, req *chunk.Chunk) error {
 	us.memBuf.RLock()
 	defer us.memBuf.RUnlock()
+
+	// Assume req.Capacity() > 0 after GrowAndReset(), if this assumption fail,
+	// the for-loop may exit without read one single row!
 	req.GrowAndReset(us.maxChunkSize)
+
 	mutableRow := chunk.MutRowFromTypes(retTypes(us))
 	for batchSize := req.Capacity(); req.NumRows() < batchSize; {
 		row, err := us.getOneRow(ctx)
diff --git a/executor/union_scan_test.go b/executor/union_scan_test.go
index 1acba89d118fb..8800e16d750dd 100644
--- a/executor/union_scan_test.go
+++ b/executor/union_scan_test.go
@@ -534,6 +534,19 @@ func TestIssue32422(t *testing.T) {
 	tk.MustExec("rollback")
 }
 
+func TestIssue36903(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t_vwvgdc")
+
+	tk.MustExec("CREATE TABLE t_vwvgdc (wkey int, pkey int NOT NULL, c_rdsfbc double DEFAULT NULL, PRIMARY KEY (`pkey`));")
+	tk.MustExec("insert into t_vwvgdc values (2, 15000, 61.75);")
+	tk.MustExec("BEGIN OPTIMISTIC;")
+	tk.MustExec("insert into t_vwvgdc (wkey, pkey, c_rdsfbc) values (155, 228000, 99.50);")
+	tk.MustQuery("select pkey from t_vwvgdc where 0 <> 0 union select pkey from t_vwvgdc;")
+}
+
 func BenchmarkUnionScanRead(b *testing.B) {
 	store := testkit.CreateMockStore(b)
 
diff --git a/util/chunk/chunk.go b/util/chunk/chunk.go
index 1dd2b4cff02e4..7270ab3d16e3c 100644
--- a/util/chunk/chunk.go
+++ b/util/chunk/chunk.go
@@ -322,7 +322,11 @@ func reCalcCapacity(c *Chunk, maxChunkSize int) int {
 	if c.NumRows() < c.capacity {
 		return c.capacity
 	}
-	return mathutil.Min(c.capacity*2, maxChunkSize)
+	newCapacity := c.capacity * 2
+	if newCapacity == 0 {
+		newCapacity = InitialCapacity
+	}
+	return mathutil.Min(newCapacity, maxChunkSize)
 }
 
 // Capacity returns the capacity of the Chunk.