executor, util: fix UnionScan Next() skip reading data when passed ch…

…unk capacity is 0 (#36961)

close #36903

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)

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)
 

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.