*: Support required rows rows for arrow decode format.

distsql/distsql_test.go

@@ -529,3 +529,284 @@ func BenchmarkDecodeToChunkWithNewChunk(b *testing.B) {
 		codec.DecodeToChunk(buffer, chk)
 	}
 }
+
+func BenchmarkDecodeToChunkWithRequestedRows_int(b *testing.B) {
+	numCols := 4
+	numRows := 1024
+
+	colTypes := make([]*types.FieldType, numCols)
+	for i := 0; i < numCols; i++ {
+		colTypes[i] = &types.FieldType{Tp: mysql.TypeLonglong}
+	}
+	chk := chunk.New(colTypes, numRows, numRows)
+
+	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
+		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
+			chk.AppendInt64(colOrdinal, 123)
+		}
+	}
+
+	codec := chunk.NewCodec(colTypes)
+	buffer := codec.Encode(chk)
+
+	acodec := chunk.NewDecoder(
+		chunk.NewChunkWithCapacity(colTypes, 0),
+		colTypes)
+
+	chk.SetRequiredRows(1024, 1024)
+	chk.Reset()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		acodec.Reset(buffer)
+		for !acodec.IsFinished() {
+			for !chk.IsFull() && !acodec.IsFinished() {
+				acodec.Decode(chk)
+			}
+			chk.Reset()
+		}
+	}
+}
+
+func BenchmarkDecodeToChunkWithRequestedRows_string(b *testing.B) {
+	numCols := 4
+	numRows := 1024
+
+	colTypes := make([]*types.FieldType, numCols)
+	for i := 0; i < numCols; i++ {
+		colTypes[i] = &types.FieldType{Tp: mysql.TypeString}
+	}
+	chk := chunk.New(colTypes, numRows, numRows)
+
+	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
+		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
+			chk.AppendString(colOrdinal, "123456")
+		}
+	}
+
+	codec := chunk.NewCodec(colTypes)
+	buffer := codec.Encode(chk)
+
+	acodec := chunk.NewDecoder(
+		chunk.NewChunkWithCapacity(colTypes, 0),
+		colTypes)
+
+	chk.SetRequiredRows(1024, 1024)
+	chk.Reset()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		acodec.Reset(buffer)
+		for !acodec.IsFinished() {
+			for !chk.IsFull() && !acodec.IsFinished() {
+				acodec.Decode(chk)
+			}
+			chk.Reset()
+		}
+	}
+}
+
+func BenchmarkDecodeToChunkReuse_int(b *testing.B) {
+	numCols := 4
+	numRows := 1024
+
+	colTypes := make([]*types.FieldType, numCols)
+	for i := 0; i < numCols; i++ {
+		colTypes[i] = &types.FieldType{Tp: mysql.TypeLonglong}
+	}
+	chk := chunk.New(colTypes, numRows, numRows)
+
+	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
+		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
+			chk.AppendInt64(colOrdinal, 123)
+		}
+	}
+
+	codec := chunk.NewCodec(colTypes)
+	buffer := codec.Encode(chk)
+
+	acodec := chunk.NewDecoder(
+		chunk.NewChunkWithCapacity(colTypes, 0),
+		colTypes)
+
+	chk.SetRequiredRows(1024, 1024)
+	chk.Reset()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		acodec.Reset(buffer)
+		for !acodec.IsFinished() {
+			for !chk.IsFull() && !acodec.IsFinished() {
+				acodec.ReuseIntermChk(chk)
+			}
+			chk.Reset()
+		}
+	}
+}
+
+func BenchmarkDecodeToChunkReuse_string(b *testing.B) {
+	numCols := 4
+	numRows := 1024
+
+	colTypes := make([]*types.FieldType, numCols)
+	for i := 0; i < numCols; i++ {
+		colTypes[i] = &types.FieldType{Tp: mysql.TypeString}
+	}
+	chk := chunk.New(colTypes, numRows, numRows)
+
+	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
+		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
+			chk.AppendString(colOrdinal, "123456")
+		}
+	}
+
+	codec := chunk.NewCodec(colTypes)
+	buffer := codec.Encode(chk)
+
+	acodec := chunk.NewDecoder(
+		chunk.NewChunkWithCapacity(colTypes, 0),
+		colTypes)
+
+	chk.SetRequiredRows(1024, 1024)
+	chk.Reset()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		acodec.Reset(buffer)
+		for !acodec.IsFinished() {
+			for !chk.IsFull() && !acodec.IsFinished() {
+				acodec.ReuseIntermChk(chk)
+			}
+			chk.Reset()
+		}
+	}
+}
+
+type dataBufferGen struct {
+	times int
+	data  [][]byte
+	cur   int
+}
+
+func (g *dataBufferGen) get() []byte {
+	if g.cur == len(g.data) {
+		g.cur, g.times = 0, g.times-1
+	}
+	if g.times == 0 {
+		return nil
+	}
+	res := make([]byte, 0, 1024)
+	res = append(res[:0], g.data[g.cur]...)
+	return res
+}
+
+func populateBufferForChunk(numCols, numRows int) []byte {
+	colTypes := make([]*types.FieldType, numCols)
+	for i := 0; i < numCols; i++ {
+		colTypes[i] = &types.FieldType{Tp: mysql.TypeString}
+	}
+	chk := chunk.New(colTypes, numRows, numRows)
+
+	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
+		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
+			chk.AppendString(colOrdinal, "123456")
+		}
+	}
+	codec := chunk.NewCodec(colTypes)
+	return codec.Encode(chk)
+}
+
+func buildBigResponse(gen *dataBufferGen) {
+	gen.data = make([][]byte, 20)
+	for i := 0; i < 20; i++ {
+		gen.data[i] = make([]byte, 0, 1024)
+	}
+	betchSize := 32
+	for i := 0; i < 20; i++ {
+		gen.data[i] = populateBufferForChunk(4, betchSize)
+		if betchSize < 1024 {
+			betchSize *= 2
+		}
+	}
+}
+
+func BenchmarkDecodeToChunkRequiredRows_BigResponse(b *testing.B) {
+	datagen := &dataBufferGen{}
+	buildBigResponse(datagen)
+
+	colTypes := make([]*types.FieldType, 4)
+	for i := 0; i < 4; i++ {
+		colTypes[i] = &types.FieldType{Tp: mysql.TypeString}
+	}
+	chk := chunk.New(colTypes, 4, 1024)
+	acodec := chunk.NewDecoder(
+		chunk.NewChunkWithCapacity(colTypes, 0),
+		colTypes)
+
+	chk.SetRequiredRows(100, 1024)
+	chk.Reset()
+	datagen.times = b.N
+	datagen.cur = 0
+	b.ResetTimer()
+
+	for true {
+		if chk.IsFull() {
+			chk.Reset()
+		}
+		if acodec.IsFinished() {
+			b.StopTimer()
+			buffer := datagen.get()
+			if buffer == nil {
+				return
+			}
+			b.StartTimer()
+			acodec.Reset(buffer)
+		}
+		// If the next chunk size is greater than required rows * 0.8, reuse the memory of the next chunk and return
+		// immediately. Otherwise, splice the data to one chunk and wait the next chunk.
+		if acodec.RemainedRows() > int(float64(chk.RequiredRows())*0.8) {
+			if chk.NumRows() > 0 {
+				chk.Reset()
+				continue
+			}
+			acodec.ReuseIntermChk(chk)
+			chk.Reset()
+			datagen.cur++
+			continue
+		}
+		acodec.Decode(chk)
+		if acodec.IsFinished() {
+			datagen.cur++
+		}
+	}
+}
+
+func BenchmarkDecodeToChunk_BigResponse(b *testing.B) {
+	datagen := &dataBufferGen{}
+	buildBigResponse(datagen)
+
+	colTypes := make([]*types.FieldType, 4)
+	for i := 0; i < 4; i++ {
+		colTypes[i] = &types.FieldType{Tp: mysql.TypeString}
+	}
+	chk := chunk.New(colTypes, 4, 1024)
+	chk.SetRequiredRows(100, 1024)
+	chk.Reset()
+	datagen.times = b.N
+	datagen.cur = 0
+	b.ResetTimer()
+
+	for true {
+		b.StopTimer()
+		buffer := datagen.get()
+		if buffer == nil {
+			return
+		}
+		b.StartTimer()
+
+		codec := chunk.NewCodec(colTypes)
+		_ = codec.DecodeToChunkTest(buffer, chk)
+		datagen.cur++
+		chk.Reset()
+	}
+}

distsql/select_result.go

@@ -66,9 +66,10 @@ type selectResult struct {
 	fieldTypes []*types.FieldType
 	ctx        sessionctx.Context
 
-	selectResp     *tipb.SelectResponse
-	selectRespSize int // record the selectResp.Size() when it is initialized.
-	respChkIdx     int
+	selectResp       *tipb.SelectResponse
+	selectRespSize   int // record the selectResp.Size() when it is initialized.
+	respChkIdx       int
+	respChunkDecoder *chunk.Decoder
 
 	feedback     *statistics.QueryFeedback
 	partialCount int64 // number of partial results.
@@ -183,10 +184,39 @@ func (r *selectResult) readFromDefault(ctx context.Context, chk *chunk.Chunk) er
 }
 
 func (r *selectResult) readFromArrow(ctx context.Context, chk *chunk.Chunk) error {
-	rowBatchData := r.selectResp.Chunks[r.respChkIdx].RowsData
-	codec := chunk.NewCodec(r.fieldTypes)
-	_ = codec.DecodeToChunk(rowBatchData, chk)
-	r.respChkIdx++
+	if r.respChunkDecoder == nil {
+		r.respChunkDecoder = chunk.NewDecoder(
+			chunk.NewChunkWithCapacity(r.fieldTypes, 0),
+			r.fieldTypes,
+		)
+	}
+
+	for !chk.IsFull() {
+		if r.respChkIdx == len(r.selectResp.Chunks) {
+			err := r.getSelectResp()
+			if err != nil || r.selectResp == nil {
+				return err
+			}
+		}
+
+		if r.respChunkDecoder.IsFinished() {
+			r.respChunkDecoder.Reset(r.selectResp.Chunks[r.respChkIdx].RowsData)
+		}
+		// If the next chunk size is greater than required rows * 0.8, reuse the memory of the next chunk and return
+		// immediately. Otherwise, splice the data to one chunk and wait the next chunk.
+		if r.respChunkDecoder.RemainedRows() > int(float64(chk.RequiredRows())*0.8) {
+			if chk.NumRows() > 0 {
+				return nil
+			}
+			r.respChunkDecoder.ReuseIntermChk(chk)
+			r.respChkIdx++
+			return nil
+		}
+		r.respChunkDecoder.Decode(chk)
+		if r.respChunkDecoder.IsFinished() {
+			r.respChkIdx++
+		}
+	}
 	return nil
 }
 

expression/integration_test.go

@@ -5031,3 +5031,49 @@ func (s *testIntegrationSuite) TestNotExistFunc(c *C) {
 	c.Assert(err.Error(), Equals, "[expression:1305]FUNCTION test.timestampliteral does not exist")
 
 }
+
+func (s *testIntegrationSuite) TestDecodetoChunkReuse(c *C) {
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk.MustExec("create table chk (a int,b varchar(20))")
+	for i := 0; i < 200; i++ {
+		if i%5 == 0 {
+			tk.MustExec(fmt.Sprintf("insert chk values (NULL,NULL)"))
+			continue
+		}
+		tk.MustExec(fmt.Sprintf("insert chk values (%d,'%s')", i, strconv.Itoa(i)))
+	}
+
+	tk.Se.GetSessionVars().DistSQLScanConcurrency = 1
+	tk.MustExec("set tidb_init_chunk_size = 2")
+	tk.MustExec("set tidb_max_chunk_size = 32")
+	defer func() {
+		tk.MustExec(fmt.Sprintf("set tidb_init_chunk_size = %d", variable.DefInitChunkSize))
+		tk.MustExec(fmt.Sprintf("set tidb_max_chunk_size = %d", variable.DefMaxChunkSize))
+	}()
+	rs, err := tk.Exec("select * from chk")
+	c.Assert(err, IsNil)
+	req := rs.NewChunk()
+	var count int
+	for {
+		err = rs.Next(context.TODO(), req)
+		c.Assert(err, IsNil)
+		numRows := req.NumRows()
+		if numRows == 0 {
+			break
+		}
+		for i := 0; i < numRows; i++ {
+			if count%5 == 0 {
+				c.Assert(req.GetRow(i).IsNull(0), Equals, true)
+				c.Assert(req.GetRow(i).IsNull(1), Equals, true)
+			} else {
+				c.Assert(req.GetRow(i).IsNull(0), Equals, false)
+				c.Assert(req.GetRow(i).IsNull(1), Equals, false)
+				c.Assert(req.GetRow(i).GetInt64(0), Equals, int64(count))
+				c.Assert(req.GetRow(i).GetString(1), Equals, strconv.Itoa(count))
+			}
+			count++
+		}
+	}
+	c.Assert(count, Equals, 200)
+	rs.Close()
+}