clean code, move precompute logic to tableReaderExecutor

executor/builder.go

@@ -1862,6 +1862,7 @@ func buildNoRangeTableReader(b *executorBuilder, v *plannercore.PhysicalTableRea
 		corColInAccess: b.corColInAccess(v.TablePlans[0]),
 		plans:          v.TablePlans,
 	}
+	e.buildVirtualColumnInfo()
 	if containsLimit(dagReq.Executors) {
 		e.feedback = statistics.NewQueryFeedback(0, nil, 0, ts.Desc)
 	} else {

executor/distsql.go

@@ -536,6 +536,7 @@ func (e *IndexLookUpExecutor) buildTableReader(ctx context.Context, handles []in
 		corColInFilter: e.corColInTblSide,
 		plans:          e.tblPlans,
 	}
+	tableReaderExec.buildVirtualColumnInfo()
 	tableReader, err := e.dataReaderBuilder.buildTableReaderFromHandles(ctx, tableReaderExec, handles)
 	if err != nil {
 		logutil.Logger(ctx).Error("build table reader from handles failed", zap.Error(err))

executor/table_reader.go

@@ -21,7 +21,6 @@ import (
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/parser/model"
 	"github.com/pingcap/tidb/distsql"
-	"github.com/pingcap/tidb/expression"
 	"github.com/pingcap/tidb/kv"
 	plannercore "github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/sessionctx"
@@ -62,7 +61,7 @@ type TableReaderExecutor struct {
 	// kvRanges are only use for union scan.
 	kvRanges []kv.KeyRange
 	dagPB    *tipb.DAGRequest
-	// columns are only required by union scan.
+	// columns are only required by union scan and virtual column.
 	columns []*model.ColumnInfo
 
 	// resultHandler handles the order of the result. Since (MAXInt64, MAXUint64] stores before [0, MaxInt64] physically
@@ -81,6 +80,11 @@ type TableReaderExecutor struct {
 	corColInFilter bool
 	// corColInAccess tells whether there's correlated column in access conditions.
 	corColInAccess bool
+	// virtualColumnIndex records all the indices of virtual columns and sort them in definition
+	// to make sure we can compute the virtual column in right order.
+	virtualColumnIndex []int
+	// virtualColumnRetFieldTypes records the RetFieldTypes of virtual columns.
+	virtualColumnRetFieldTypes []*types.FieldType
 }
 
 // Open initialzes necessary variables for using this executor.
@@ -159,43 +163,22 @@ func (e *TableReaderExecutor) Next(ctx context.Context, req *chunk.Chunk) error
 		return err
 	}
 
-	// If e.schema contains virtual column, compute the virtual column and put them into the chunk
-	// TODO: move them to other place so that we don't need to compute every Next()
-	virColTypes := make([]*types.FieldType, 0)
-	virColDef := make([]*expression.Column, 0)
-	virColIndex := make([]int, 0)
-	virColOrder := make([]int, 0)
-	for i, col := range e.schema.Columns {
-		if col.VirtualExpr != nil {
-			virColTypes = append(virColTypes, col.RetType)
-			virColDef = append(virColDef, col)
-			virColIndex = append(virColIndex, i)
-			virColOrder = append(virColOrder, len(virColOrder))
-		}
-	}
-	if len(virColOrder) > 0 {
-		sort.Slice(virColOrder, func(i, j int) bool {
-			return model.FindColumnInfo(e.columns, virColDef[i].OrigColName.String()).Offset <
-				model.FindColumnInfo(e.columns, virColDef[j].OrigColName.String()).Offset
-		})
-	}
-
-	virCols := chunk.NewChunkWithCapacity(virColTypes, req.Capacity())
+	virCols := chunk.NewChunkWithCapacity(e.virtualColumnRetFieldTypes, req.Capacity())
 	iter := chunk.NewIterator4Chunk(req)
 
-	for _, j := range virColOrder {
+	for i, idx := range e.virtualColumnIndex {
 		for row := iter.Begin(); row != iter.End(); row = iter.Next() {
-			datum, err := virColDef[j].EvalVirtualColumn(row)
+			datum, err := e.schema.Columns[idx].EvalVirtualColumn(row)
 			if err != nil {
 				return err
 			}
-			castDatum, err := table.CastValue(e.ctx, datum, e.columns[virColIndex[j]])
+			castDatum, err := table.CastValue(e.ctx, datum, e.columns[idx])
 			if err != nil {
 				return err
 			}
-			virCols.AppendDatum(j, &castDatum)
+			virCols.AppendDatum(i, &castDatum)
 		}
-		req.SetCol(virColIndex[j], virCols.Column(j))
+		req.SetCol(idx, virCols.Column(i))
 	}
 
 	return nil
@@ -239,6 +222,26 @@ func (e *TableReaderExecutor) buildResp(ctx context.Context, ranges []*ranger.Ra
 	return result, nil
 }
 
+// buildVirtualColumnInfo saves virtual column indices and sort them in definition order
+func (e *TableReaderExecutor) buildVirtualColumnInfo() {
+	e.virtualColumnIndex = make([]int, 0)
+	for i, col := range e.schema.Columns {
+		if col.VirtualExpr != nil {
+			e.virtualColumnIndex = append(e.virtualColumnIndex, i)
+		}
+	}
+	sort.Slice(e.virtualColumnIndex, func(i, j int) bool {
+		return model.FindColumnInfo(e.columns, e.schema.Columns[e.virtualColumnIndex[i]].OrigColName.String()).Offset <
+			model.FindColumnInfo(e.columns, e.schema.Columns[e.virtualColumnIndex[j]].OrigColName.String()).Offset
+	})
+	if len(e.virtualColumnIndex) > 0 {
+		e.virtualColumnRetFieldTypes = make([]*types.FieldType, 0)
+		for _, idx := range e.virtualColumnIndex {
+			e.virtualColumnRetFieldTypes = append(e.virtualColumnRetFieldTypes, e.schema.Columns[idx].RetType)
+		}
+	}
+}
+
 type tableResultHandler struct {
 	// If the pk is unsigned and we have KeepOrder=true and want ascending order,
 	// `optionalResult` will handles the request whose range is in signed int range, and