diff --git a/planner/core/planbuilder.go b/planner/core/planbuilder.go index c1c4eb65e8..3f8fa33bb7 100644 --- a/planner/core/planbuilder.go +++ b/planner/core/planbuilder.go @@ -43,6 +43,7 @@ import ( "github.com/pingcap/tidb/util/chunk" "github.com/pingcap/tidb/util/logutil" "github.com/pingcap/tidb/util/ranger" + "github.com/pingcap/tidb/util/set" "go.uber.org/zap" ) @@ -743,14 +744,14 @@ func (b *PlanBuilder) buildAdmin(ctx context.Context, as *ast.AdminStmt) (Plan, } // getGenExprs gets generated expressions map. -func (b *PlanBuilder) getGenExprs(ctx context.Context, dbName model.CIStr, tbl table.Table, idx *model.IndexInfo) ( +func (b *PlanBuilder) getGenExprs(ctx context.Context, dbName model.CIStr, tbl table.Table, idx *model.IndexInfo, exprCols *expression.Schema) ( map[model.TableColumnID]expression.Expression, error) { tblInfo := tbl.Meta() genExprsMap := make(map[model.TableColumnID]expression.Expression) exprs := make([]expression.Expression, 0, len(tbl.Cols())) genExprIdxs := make([]model.TableColumnID, len(tbl.Cols())) mockTablePlan := LogicalTableDual{}.Init(b.ctx) - mockTablePlan.SetSchema(expression.TableInfo2SchemaWithDBName(b.ctx, dbName, tblInfo)) + mockTablePlan.SetSchema(exprCols) for i, colExpr := range mockTablePlan.Schema().Columns { col := tbl.Cols()[i] var expr expression.Expression @@ -787,60 +788,68 @@ func (b *PlanBuilder) getGenExprs(ctx context.Context, dbName model.CIStr, tbl t return genExprsMap, nil } -func (b *PlanBuilder) buildPhysicalIndexLookUpReader(ctx context.Context, dbName model.CIStr, tbl table.Table, idx *model.IndexInfo, id int) (Plan, error) { - genExprsMap, err := b.getGenExprs(ctx, dbName, tbl, idx) - if err != nil { - return nil, errors.Trace(err) +func findColumnInfoByID(colInfos []*model.ColumnInfo, id int64) *model.ColumnInfo { + for _, info := range colInfos { + if info.ID == id { + return info + } } + return nil +} +func (b *PlanBuilder) buildPhysicalIndexLookUpReader(ctx context.Context, dbName model.CIStr, tbl table.Table, idx *model.IndexInfo, id int) (Plan, error) { // Get generated columns. var genCols []*expression.Column pkOffset := -1 tblInfo := tbl.Meta() - colsMap := make(map[int64]struct{}) + colsMap := set.NewInt64Set() schema := expression.NewSchema(make([]*expression.Column, 0, len(idx.Columns))...) idxReaderCols := make([]*model.ColumnInfo, 0, len(idx.Columns)) tblReaderCols := make([]*model.ColumnInfo, 0, len(tbl.Cols())) + fullExprCols := expression.TableInfo2SchemaWithDBName(b.ctx, dbName, tblInfo) + genExprsMap, err := b.getGenExprs(ctx, dbName, tbl, idx, fullExprCols) + if err != nil { + return nil, err + } for _, idxCol := range idx.Columns { - for _, col := range tblInfo.Columns { + for i, col := range tblInfo.Columns { if idxCol.Name.L == col.Name.L { idxReaderCols = append(idxReaderCols, col) tblReaderCols = append(tblReaderCols, col) - schema.Append(&expression.Column{ - ColName: col.Name, - UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(), - RetType: &col.FieldType}) - colsMap[col.ID] = struct{}{} + schema.Append(fullExprCols.Columns[i]) + colsMap.Insert(col.ID) if mysql.HasPriKeyFlag(col.Flag) { pkOffset = len(tblReaderCols) - 1 } - } - genColumnID := model.TableColumnID{TableID: tblInfo.ID, ColumnID: col.ID} - if expr, ok := genExprsMap[genColumnID]; ok { - cols := expression.ExtractColumns(expr) - genCols = append(genCols, cols...) + genColumnID := model.TableColumnID{TableID: tblInfo.ID, ColumnID: col.ID} + if expr, ok := genExprsMap[genColumnID]; ok { + cols := expression.ExtractColumns(expr) + genCols = append(genCols, cols...) + } } } } // Add generated columns to tblSchema and tblReaderCols. tblSchema := schema.Clone() for _, col := range genCols { - if _, ok := colsMap[col.ID]; !ok { - c := table.FindCol(tbl.Cols(), col.ColName.O) - if c != nil { - col.Index = len(tblReaderCols) - tblReaderCols = append(tblReaderCols, c.ColumnInfo) - tblSchema.Append(&expression.Column{ - ColName: c.Name, - UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(), - RetType: &c.FieldType}) - colsMap[c.ID] = struct{}{} - if mysql.HasPriKeyFlag(c.Flag) { + if !colsMap.Exist(col.ID) { + info := findColumnInfoByID(tblInfo.Columns, col.ID) + if info != nil { + tblReaderCols = append(tblReaderCols, info) + tblSchema.Append(col) + colsMap.Insert(col.ID) + if mysql.HasPriKeyFlag(col.RetType.Flag) { pkOffset = len(tblReaderCols) - 1 } } } } + for k, expr := range genExprsMap { + genExprsMap[k], err = expr.ResolveIndices(tblSchema) + if err != nil { + return nil, err + } + } if !tbl.Meta().PKIsHandle || pkOffset == -1 { tblReaderCols = append(tblReaderCols, model.NewExtraHandleColInfo()) handleCol := &expression.Column{