*: Support required rows rows for arrow decode format.

distsql/distsql_test.go

@@ -25,7 +25,7 @@ import (
 	"github.com/pingcap/parser/charset"
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/kv"
-	"github.com/pingcap/tidb/sessionctx/stmtctx"
+	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/types"
@@ -127,6 +127,7 @@ func (s *testSuite) TestSelectMemTracker(c *C) {
 }
 
 func (s *testSuite) TestSelectNormalChunkSize(c *C) {
+	s.sctx.GetSessionVars().EnableArrow = false
 	response, colTypes := s.createSelectNormal(100, 1000000, c, nil)
 	response.Fetch(context.TODO())
 	s.testChunkSize(response, colTypes, c)
@@ -288,6 +289,7 @@ func (s *testSuite) testChunkSize(response SelectResult, colTypes []*types.Field
 }
 
 func (s *testSuite) TestAnalyze(c *C) {
+	s.sctx.GetSessionVars().EnableArrow = false
 	request, err := (&RequestBuilder{}).SetKeyRanges(nil).
 		SetAnalyzeRequest(&tipb.AnalyzeReq{}).
 		SetKeepOrder(true).
@@ -313,6 +315,7 @@ func (s *testSuite) TestAnalyze(c *C) {
 }
 
 func (s *testSuite) TestChecksum(c *C) {
+	s.sctx.GetSessionVars().EnableArrow = false
 	request, err := (&RequestBuilder{}).SetKeyRanges(nil).
 		SetChecksumRequest(&tipb.ChecksumRequest{}).
 		Build()
@@ -342,6 +345,7 @@ type mockResponse struct {
 	count int
 	total int
 	batch int
+	ctx   sessionctx.Context
 	sync.Mutex
 }
 
@@ -365,20 +369,50 @@ func (resp *mockResponse) Next(ctx context.Context) (kv.ResultSubset, error) {
 	numRows := mathutil.Min(resp.batch, resp.total-resp.count)
 	resp.count += numRows
 
-	datum := types.NewIntDatum(1)
-	bytes := make([]byte, 0, 100)
-	bytes, _ = codec.EncodeValue(nil, bytes, datum, datum, datum, datum)
-	chunks := make([]tipb.Chunk, numRows)
-	for i := range chunks {
-		chkData := make([]byte, len(bytes))
-		copy(chkData, bytes)
-		chunks[i] = tipb.Chunk{RowsData: chkData}
+	var chunks []tipb.Chunk
+	if !enableTypeArrow(resp.ctx) {
+		datum := types.NewIntDatum(1)
+		bytes := make([]byte, 0, 100)
+		bytes, _ = codec.EncodeValue(nil, bytes, datum, datum, datum, datum)
+		chunks = make([]tipb.Chunk, numRows)
+		for i := range chunks {
+			chkData := make([]byte, len(bytes))
+			copy(chkData, bytes)
+			chunks[i] = tipb.Chunk{RowsData: chkData}
+		}
+	} else {
+		chunks = make([]tipb.Chunk, 0)
+		for numRows > 0 {
+			rows := mathutil.Min(numRows, 1024)
+			numRows -= rows
+
+			colTypes := make([]*types.FieldType, 4)
+			for i := 0; i < 4; i++ {
+				colTypes[i] = &types.FieldType{Tp: mysql.TypeLonglong}
+			}
+			chk := chunk.New(colTypes, numRows, numRows)
+
+			for rowOrdinal := 0; rowOrdinal < rows; rowOrdinal++ {
+				for colOrdinal := 0; colOrdinal < 4; colOrdinal++ {
+					chk.AppendInt64(colOrdinal, 123)
+				}
+			}
+
+			codec := chunk.NewCodec(colTypes)
+			buffer := codec.Encode(chk)
+			chunks = append(chunks, tipb.Chunk{RowsData: buffer})
+		}
 	}
 
 	respPB := &tipb.SelectResponse{
 		Chunks:       chunks,
 		OutputCounts: []int64{1},
 	}
+	if enableTypeArrow(resp.ctx) {
+		respPB.EncodeType = tipb.EncodeType_TypeArrow
+	} else {
+		respPB.EncodeType = tipb.EncodeType_TypeDefault
+	}
 	respBytes, err := respPB.Marshal()
 	if err != nil {
 		panic(err)
@@ -407,125 +441,85 @@ func (r *mockResultSubset) MemSize() int64 { return int64(cap(r.data)) }
 // RespTime implements kv.ResultSubset interface.
 func (r *mockResultSubset) RespTime() time.Duration { return 0 }
 
-func populateBuffer() []byte {
-	numCols := 4
-	numRows := 1024
-	buffer := make([]byte, 0, 1024)
-	sc := &stmtctx.StatementContext{TimeZone: time.Local}
-
-	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
-		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
-			buffer, _ = codec.EncodeValue(sc, buffer, types.NewIntDatum(123))
-		}
-	}
-
-	return buffer
-}
-
-func mockReadRowsData(buffer []byte, colTypes []*types.FieldType, chk *chunk.Chunk) (err error) {
-	chk.Reset()
-	numCols := 4
-	numRows := 1024
-
-	decoder := codec.NewDecoder(chk, time.Local)
-	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
-		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
-			buffer, err = decoder.DecodeOne(buffer, colOrdinal, colTypes[colOrdinal])
-			if err != nil {
-				return err
-			}
-		}
-	}
-
-	return nil
-}
-
-func BenchmarkReadRowsData(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)
-
-	buffer := populateBuffer()
-
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		mockReadRowsData(buffer, colTypes, chk)
-	}
-}
-
-func BenchmarkDecodeToChunk(b *testing.B) {
-	numCols := 4
-	numRows := 1024
+func createSelectNormal(batch, totalRows int, ctx sessionctx.Context) (*selectResult, []*types.FieldType) {
+	request, _ := (&RequestBuilder{}).SetKeyRanges(nil).
+		SetDAGRequest(&tipb.DAGRequest{}).
+		SetDesc(false).
+		SetKeepOrder(false).
+		SetFromSessionVars(variable.NewSessionVars()).
+		SetMemTracker(memory.NewTracker(stringutil.StringerStr("testSuite.createSelectNormal"),
+			ctx.GetSessionVars().MemQuotaDistSQL)).
+		Build()
 
-	colTypes := make([]*types.FieldType, numCols)
-	for i := 0; i < numCols; i++ {
-		colTypes[i] = &types.FieldType{Tp: mysql.TypeLonglong}
+	/// 4 int64 types.
+	colTypes := []*types.FieldType{
+		{
+			Tp:      mysql.TypeLonglong,
+			Flen:    mysql.MaxIntWidth,
+			Decimal: 0,
+			Flag:    mysql.BinaryFlag,
+			Charset: charset.CharsetBin,
+			Collate: charset.CollationBin,
+		},
 	}
-	chk := chunk.New(colTypes, numRows, numRows)
+	colTypes = append(colTypes, colTypes[0])
+	colTypes = append(colTypes, colTypes[0])
+	colTypes = append(colTypes, colTypes[0])
 
-	for rowOrdinal := 0; rowOrdinal < numRows; rowOrdinal++ {
-		for colOrdinal := 0; colOrdinal < numCols; colOrdinal++ {
-			chk.AppendInt64(colOrdinal, 123)
-		}
-	}
+	// Test Next.
+	var response SelectResult
+	response, _ = Select(context.TODO(), ctx, request, colTypes, statistics.NewQueryFeedback(0, nil, 0, false))
 
-	codec := chunk.NewCodec(colTypes)
-	buffer := codec.Encode(chk)
+	result, _ := response.(*selectResult)
+	resp, _ := result.resp.(*mockResponse)
+	resp.total = totalRows
+	resp.batch = batch
 
-	b.ResetTimer()
-	for i := 0; i < b.N; i++ {
-		codec.DecodeToChunk(buffer, chk)
-	}
+	return result, colTypes
 }
 
-func BenchmarkReadRowsDataWithNewChunk(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}
-	}
-	buffer := populateBuffer()
-
-	b.ResetTimer()
+func BenchmarkSelectResponseChunk_BigResponse(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		b.StopTimer()
-		chk := chunk.New(colTypes, numRows, numRows)
+		s := &testSuite{}
+		s.SetUpSuite(nil)
+		selectResult, colTypes := createSelectNormal(4000, 20000, s.sctx)
+		selectResult.Fetch(context.TODO())
+		chk := chunk.NewChunkWithCapacity(colTypes, 1024)
 		b.StartTimer()
-		mockReadRowsData(buffer, colTypes, chk)
-	}
-}
-
-func BenchmarkDecodeToChunkWithNewChunk(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)
+		for true {
+			err := selectResult.Next(context.TODO(), chk)
+			if err != nil {
+				panic(err)
+			}
+			if chk.NumRows() == 0 {
+				break
+			}
+			chk.Reset()
 		}
+		s.TearDownSuite(nil)
 	}
+}
 
-	codec := chunk.NewCodec(colTypes)
-	buffer := codec.Encode(chk)
-
-	b.ResetTimer()
+func BenchmarkSelectResponseChunk_SmallResponse(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		b.StopTimer()
-		chk := chunk.New(colTypes, numRows, numRows)
+		s := &testSuite{}
+		s.SetUpSuite(nil)
+		selectResult, colTypes := createSelectNormal(32, 3200, s.sctx)
+		selectResult.Fetch(context.TODO())
+		chk := chunk.NewChunkWithCapacity(colTypes, 1024)
 		b.StartTimer()
-		codec.DecodeToChunk(buffer, chk)
+		for true {
+			err := selectResult.Next(context.TODO(), chk)
+			if err != nil {
+				panic(err)
+			}
+			if chk.NumRows() == 0 {
+				break
+			}
+			chk.Reset()
+		}
+		s.TearDownSuite(nil)
 	}
 }

distsql/request_builder_test.go

@@ -60,6 +60,7 @@ func (s *testSuite) SetUpSuite(c *C) {
 	ctx.Store = &mock.Store{
 		Client: &mock.Client{
 			MockResponse: &mockResponse{
+				ctx:   ctx,
 				batch: 1,
 				total: 2,
 			},
@@ -77,6 +78,7 @@ func (s *testSuite) SetUpTest(c *C) {
 	store := ctx.Store.(*mock.Store)
 	store.Client = &mock.Client{
 		MockResponse: &mockResponse{
+			ctx:   ctx,
 			batch: 1,
 			total: 2,
 		},

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
 }