*: encode the returned data using the Chunk format in mocktikv

config/config.go

@@ -287,6 +287,8 @@ type TiKVClient struct {
 	MaxBatchWaitTime time.Duration `toml:"max-batch-wait-time" json:"max-batch-wait-time"`
 	// BatchWaitSize is the max wait size for batch.
 	BatchWaitSize uint `toml:"batch-wait-size" json:"batch-wait-size"`
+	// EnableArrow indicate the data encode in arrow format.
+	EnableArrow bool `toml:"enable-arrow" json:"enable-arrow"`
 }
 
 // Binlog is the config for binlog.
@@ -411,6 +413,8 @@ var defaultConf = Config{
 		OverloadThreshold: 200,
 		MaxBatchWaitTime:  0,
 		BatchWaitSize:     8,
+
+		EnableArrow: true,
 	},
 	Binlog: Binlog{
 		WriteTimeout: "15s",

config/config.toml.example

@@ -285,6 +285,9 @@ max-batch-wait-time = 0
 # Batch wait size, to avoid waiting too long.
 batch-wait-size = 8
 
+# Enable chunk encoded data.
+enable-arrow = true
+
 [txn-local-latches]
 # Enable local latches for transactions. Enable it when
 # there are lots of conflicts between transactions.

ddl/reorg.go

@@ -220,7 +220,7 @@ func constructLimitPB(count uint64) *tipb.Executor {
 	return &tipb.Executor{Tp: tipb.ExecType_TypeLimit, Limit: limitExec}
 }
 
-func buildDescTableScanDAG(startTS uint64, tbl table.PhysicalTable, columns []*model.ColumnInfo, limit uint64) (*tipb.DAGRequest, error) {
+func buildDescTableScanDAG(ctx sessionctx.Context, startTS uint64, tbl table.PhysicalTable, columns []*model.ColumnInfo, limit uint64) (*tipb.DAGRequest, error) {
 	dagReq := &tipb.DAGRequest{}
 	dagReq.StartTs = startTS
 	_, timeZoneOffset := time.Now().In(time.UTC).Zone()
@@ -234,6 +234,7 @@ func buildDescTableScanDAG(startTS uint64, tbl table.PhysicalTable, columns []*m
 	tblScanExec := constructDescTableScanPB(tbl.GetPhysicalID(), pbColumnInfos)
 	dagReq.Executors = append(dagReq.Executors, tblScanExec)
 	dagReq.Executors = append(dagReq.Executors, constructLimitPB(limit))
+	distsql.SetEncodeType(ctx, dagReq)
 	return dagReq, nil
 }
 
@@ -247,7 +248,8 @@ func getColumnsTypes(columns []*model.ColumnInfo) []*types.FieldType {
 
 // buildDescTableScan builds a desc table scan upon tblInfo.
 func (d *ddlCtx) buildDescTableScan(ctx context.Context, startTS uint64, tbl table.PhysicalTable, columns []*model.ColumnInfo, limit uint64) (distsql.SelectResult, error) {
-	dagPB, err := buildDescTableScanDAG(startTS, tbl, columns, limit)
+	sctx := newContext(d.store)
+	dagPB, err := buildDescTableScanDAG(sctx, startTS, tbl, columns, limit)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -266,7 +268,6 @@ func (d *ddlCtx) buildDescTableScan(ctx context.Context, startTS uint64, tbl tab
 		return nil, errors.Trace(err)
 	}
 
-	sctx := newContext(d.store)
 	result, err := distsql.Select(ctx, sctx, kvReq, getColumnsTypes(columns), statistics.NewQueryFeedback(0, nil, 0, false))
 	if err != nil {
 		return nil, errors.Trace(err)

distsql/distsql.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tipb/go-tipb"
 )
 
 // XAPI error codes.
@@ -74,6 +75,10 @@ func Select(ctx context.Context, sctx sessionctx.Context, kvReq *kv.Request, fie
 			feedback:   fb,
 		}, nil
 	}
+	encodetype := tipb.EncodeType_TypeDefault
+	if useChunkIPC(sctx) {
+		encodetype = tipb.EncodeType_TypeArrow
+	}
 	return &selectResult{
 		label:      "dag",
 		resp:       resp,
@@ -85,6 +90,7 @@ func Select(ctx context.Context, sctx sessionctx.Context, kvReq *kv.Request, fie
 		feedback:   fb,
 		sqlType:    label,
 		memTracker: kvReq.MemTracker,
+		encodeType: encodetype,
 	}, nil
 }
 
@@ -115,12 +121,13 @@ func Analyze(ctx context.Context, client kv.Client, kvReq *kv.Request, vars *kv.
 		label = metrics.LblInternal
 	}
 	result := &selectResult{
-		label:    "analyze",
-		resp:     resp,
-		results:  make(chan resultWithErr, kvReq.Concurrency),
-		closed:   make(chan struct{}),
-		feedback: statistics.NewQueryFeedback(0, nil, 0, false),
-		sqlType:  label,
+		label:      "analyze",
+		resp:       resp,
+		results:    make(chan resultWithErr, kvReq.Concurrency),
+		closed:     make(chan struct{}),
+		feedback:   statistics.NewQueryFeedback(0, nil, 0, false),
+		sqlType:    label,
+		encodeType: tipb.EncodeType_TypeDefault,
 	}
 	return result, nil
 }
@@ -132,12 +139,35 @@ func Checksum(ctx context.Context, client kv.Client, kvReq *kv.Request, vars *kv
 		return nil, errors.New("client returns nil response")
 	}
 	result := &selectResult{
-		label:    "checksum",
-		resp:     resp,
-		results:  make(chan resultWithErr, kvReq.Concurrency),
-		closed:   make(chan struct{}),
-		feedback: statistics.NewQueryFeedback(0, nil, 0, false),
-		sqlType:  metrics.LblGeneral,
+		label:      "checksum",
+		resp:       resp,
+		results:    make(chan resultWithErr, kvReq.Concurrency),
+		closed:     make(chan struct{}),
+		feedback:   statistics.NewQueryFeedback(0, nil, 0, false),
+		sqlType:    metrics.LblGeneral,
+		encodeType: tipb.EncodeType_TypeDefault,
 	}
 	return result, nil
 }
+
+// SetEncodeType sets the encoding method for the DAGRequest. The supported encoding
+// methods are:
+// 1. TypeArrow: the result is encoded using the Chunk format, refer util/chunk/chunk.go
+// 2. TypeDefault: the result is encoded row by row
+func SetEncodeType(ctx sessionctx.Context, dagReq *tipb.DAGRequest) {
+	if useChunkIPC(ctx) {
+		dagReq.EncodeType = tipb.EncodeType_TypeArrow
+	} else {
+		dagReq.EncodeType = tipb.EncodeType_TypeDefault
+	}
+}
+
+func useChunkIPC(ctx sessionctx.Context) bool {
+	if !ctx.GetSessionVars().EnableArrow {
+		return false
+	}
+	if ctx.GetSessionVars().EnableStreaming {
+		return false
+	}
+	return true
+}

distsql/distsql_test.go

@@ -482,3 +482,50 @@ func BenchmarkDecodeToChunk(b *testing.B) {
 		codec.DecodeToChunk(buffer, chk)
 	}
 }
+
+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()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		chk := chunk.New(colTypes, numRows, numRows)
+		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)
+		}
+	}
+
+	codec := chunk.NewCodec(colTypes)
+	buffer := codec.Encode(chk)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		chk := chunk.New(colTypes, numRows, numRows)
+		b.StartTimer()
+		codec.DecodeToChunk(buffer, chk)
+	}
+}

distsql/select_result.go

@@ -73,6 +73,7 @@ type selectResult struct {
 	feedback     *statistics.QueryFeedback
 	partialCount int64 // number of partial results.
 	sqlType      string
+	encodeType   tipb.EncodeType
 
 	// copPlanIDs contains all copTasks' planIDs,
 	// which help to collect copTasks' runtime stats.
@@ -145,8 +146,30 @@ func (r *selectResult) NextRaw(ctx context.Context) (data []byte, err error) {
 // Next reads data to the chunk.
 func (r *selectResult) Next(ctx context.Context, chk *chunk.Chunk) error {
 	chk.Reset()
+	// Check the returned data is default/arrow format.
+	if r.selectResp == nil || (len(r.selectResp.RowBatchData) == 0 && r.respChkIdx == len(r.selectResp.Chunks)) {
+		err := r.getSelectResp()
+		if err != nil || r.selectResp == nil {
+			return err
+		}
+		// TODO(Shenghui Wu): add metrics
+		if len(r.selectResp.RowBatchData) == 0 {
+			r.encodeType = tipb.EncodeType_TypeDefault
+		}
+	}
+
+	switch r.encodeType {
+	case tipb.EncodeType_TypeDefault:
+		return r.readFromDefault(ctx, chk)
+	case tipb.EncodeType_TypeArrow:
+		return r.readFromArrow(ctx, chk)
+	}
+	return errors.Errorf("unsupported encode type:%v", r.encodeType)
+}
+
+func (r *selectResult) readFromDefault(ctx context.Context, chk *chunk.Chunk) error {
 	for !chk.IsFull() {
-		if r.selectResp == nil || r.respChkIdx == len(r.selectResp.Chunks) {
+		if r.respChkIdx == len(r.selectResp.Chunks) {
 			err := r.getSelectResp()
 			if err != nil || r.selectResp == nil {
 				return err
@@ -163,6 +186,14 @@ func (r *selectResult) Next(ctx context.Context, chk *chunk.Chunk) error {
 	return nil
 }
 
+func (r *selectResult) readFromArrow(ctx context.Context, chk *chunk.Chunk) error {
+	rowBatchData := r.selectResp.RowBatchData
+	codec := chunk.NewCodec(r.fieldTypes)
+	remained := codec.DecodeToChunk(rowBatchData, chk)
+	r.selectResp.RowBatchData = remained
+	return nil
+}
+
 func (r *selectResult) getSelectResp() error {
 	r.respChkIdx = 0
 	for {
@@ -196,7 +227,7 @@ func (r *selectResult) getSelectResp() error {
 		r.feedback.Update(re.result.GetStartKey(), r.selectResp.OutputCounts)
 		r.partialCount++
 		sc.MergeExecDetails(re.result.GetExecDetails(), nil)
-		if len(r.selectResp.Chunks) == 0 {
+		if len(r.selectResp.Chunks) == 0 && len(r.selectResp.RowBatchData) == 0 {
 			continue
 		}
 		return nil

executor/admin.go

@@ -146,6 +146,7 @@ func (e *CheckIndexRangeExec) buildDAGPB() (*tipb.DAGRequest, error) {
 	if err != nil {
 		return nil, err
 	}
+	distsql.SetEncodeType(e.ctx, dagReq)
 	return dagReq, nil
 }
 
@@ -249,7 +250,7 @@ func (e *RecoverIndexExec) buildDAGPB(txn kv.Transaction, limitCnt uint64) (*tip
 
 	limitExec := e.constructLimitPB(limitCnt)
 	dagReq.Executors = append(dagReq.Executors, limitExec)
-
+	distsql.SetEncodeType(e.ctx, dagReq)
 	return dagReq, nil
 }
 
@@ -684,7 +685,7 @@ func (e *CleanupIndexExec) buildIdxDAGPB(txn kv.Transaction) (*tipb.DAGRequest,
 
 	limitExec := e.constructLimitPB()
 	dagReq.Executors = append(dagReq.Executors, limitExec)
-
+	distsql.SetEncodeType(e.ctx, dagReq)
 	return dagReq, nil
 }
 

executor/builder.go

@@ -1699,6 +1699,8 @@ func (b *executorBuilder) constructDAGReq(plans []plannercore.PhysicalPlan) (dag
 	sc := b.ctx.GetSessionVars().StmtCtx
 	dagReq.Flags = sc.PushDownFlags()
 	dagReq.Executors, streaming, err = constructDistExec(b.ctx, plans)
+
+	distsql.SetEncodeType(b.ctx, dagReq)
 	return dagReq, streaming, err
 }
 

executor/table_readers_required_rows_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/distsql"
 	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/table/tables"
@@ -128,7 +129,11 @@ func buildTableReader(sctx sessionctx.Context) Executor {
 
 func buildMockDAGRequest(sctx sessionctx.Context) *tipb.DAGRequest {
 	builder := newExecutorBuilder(sctx, nil)
-	req, _, err := builder.constructDAGReq(nil)
+	req, _, err := builder.constructDAGReq([]core.PhysicalPlan{&core.PhysicalTableScan{
+		Columns: []*model.ColumnInfo{},
+		Table:   &model.TableInfo{ID: 12345, PKIsHandle: false},
+		Desc:    false,
+	}})
 	if err != nil {
 		panic(err)
 	}

expression/distsql_builtin.go

@@ -28,7 +28,8 @@ import (
 	"github.com/pingcap/tipb/go-tipb"
 )
 
-func pbTypeToFieldType(tp *tipb.FieldType) *types.FieldType {
+// PbTypeToFieldType converts tipb.FieldType to FieldType
+func PbTypeToFieldType(tp *tipb.FieldType) *types.FieldType {
 	return &types.FieldType{
 		Tp:      byte(tp.Tp),
 		Flag:    uint(tp.Flag),
@@ -40,7 +41,7 @@ func pbTypeToFieldType(tp *tipb.FieldType) *types.FieldType {
 }
 
 func getSignatureByPB(ctx sessionctx.Context, sigCode tipb.ScalarFuncSig, tp *tipb.FieldType, args []Expression) (f builtinFunc, e error) {
-	fieldTp := pbTypeToFieldType(tp)
+	fieldTp := PbTypeToFieldType(tp)
 	base := newBaseBuiltinFunc(ctx, args)
 	base.tp = fieldTp
 	switch sigCode {
@@ -551,7 +552,7 @@ func PBToExpr(expr *tipb.Expr, tps []*types.FieldType, sc *stmtctx.StatementCont
 }
 
 func convertTime(data []byte, ftPB *tipb.FieldType, tz *time.Location) (*Constant, error) {
-	ft := pbTypeToFieldType(ftPB)
+	ft := PbTypeToFieldType(ftPB)
 	_, v, err := codec.DecodeUint(data)
 	if err != nil {
 		return nil, err

session/session_test.go

@@ -2018,7 +2018,7 @@ func (s *testSchemaSuite) TestTableReaderChunk(c *C) {
 	}
 	c.Assert(count, Equals, 100)
 	// FIXME: revert this result to new group value after distsql can handle initChunkSize.
-	c.Assert(numChunks, Equals, 1)
+	c.Assert(numChunks, Equals, 10)
 	rs.Close()
 }
 

sessionctx/variable/session.go

@@ -352,6 +352,9 @@ type SessionVars struct {
 	// TODO: remove this after tidb-server configuration "enable-streaming' removed.
 	EnableStreaming bool
 
+	// EnableArrow indicates whether the coprocessor request can use arrow API.
+	EnableArrow bool
+
 	writeStmtBufs WriteStmtBufs
 
 	// L2CacheSize indicates the size of CPU L2 cache, using byte as unit.
@@ -525,6 +528,14 @@ func NewSessionVars() *SessionVars {
 		enableStreaming = "0"
 	}
 	terror.Log(vars.SetSystemVar(TiDBEnableStreaming, enableStreaming))
+
+	var enableArrow string
+	if config.GetGlobalConfig().TiKVClient.EnableArrow {
+		enableArrow = "1"
+	} else {
+		enableArrow = "0"
+	}
+	terror.Log(vars.SetSystemVar(TiDBEnableArrow, enableArrow))
 	return vars
 }
 
@@ -848,6 +859,8 @@ func (s *SessionVars) SetSystemVar(name string, val string) error {
 		s.DisableTxnAutoRetry = TiDBOptOn(val)
 	case TiDBEnableStreaming:
 		s.EnableStreaming = TiDBOptOn(val)
+	case TiDBEnableArrow:
+		s.EnableArrow = TiDBOptOn(val)
 	case TiDBEnableCascadesPlanner:
 		s.EnableCascadesPlanner = TiDBOptOn(val)
 	case TiDBOptimizerSelectivityLevel: