update

executor/index_merge_reader.go

@@ -938,18 +938,16 @@ func (w *indexMergeProcessWorker) NewHandleHeap(taskMap map[int][]*indexMergeTab
 }
 
 // pruneTableWorkerTaskIdxRows prune idxRows and only keep columns that will be used in byItems.
+// e.g. the common handle is (`b`, `c`) and order by with column `c`, we should make column `c` at the first.
 func (w *indexMergeProcessWorker) pruneTableWorkerTaskIdxRows(task *indexMergeTableTask) {
 	// IndexScan no need to prune retChk, Columns required by byItems are always first.
 	if plan, ok := w.indexMerge.partialPlans[task.partialPlanID][0].(*plannercore.PhysicalTableScan); ok {
 		prune := make([]int, 0, len(w.indexMerge.byItems))
 		for _, item := range plan.ByItems {
 			c, _ := item.Expr.(*expression.Column)
-			for i, col := range plan.Schema().Columns {
-				if c.ID == col.ID {
-					prune = append(prune, i)
-					break
-				}
-			}
+			idx := plan.Schema().ColumnIndex(c)
+			// couldn't equals to -1 here, if idx == -1, just let it panic
+			prune = append(prune, idx)
 		}
 		task.idxRows = task.idxRows.Prune(prune)
 	}

planner/core/task.go

@@ -1214,23 +1214,19 @@ func (p *PhysicalTopN) pushPartialTopNDownToCop(copTsk *copTask) (task, bool) {
 			if tblInfo.GetPartitionInfo() == nil {
 				if tblInfo.PKIsHandle {
 					pk := tblInfo.GetPkColInfo()
-					for _, c := range tblScan.tblCols {
-						if c.ID == pk.ID {
-							tblScan.HandleCols = NewIntHandleCols(c)
-							clonedTblScan.(*PhysicalTableScan).schema.Append(c)
-							clonedTblScan.(*PhysicalTableScan).Columns = append(clonedTblScan.(*PhysicalTableScan).Columns, c.ToInfo())
-							break
-						}
-					}
+					col := expression.ColInfo2Col(tblScan.tblCols, pk)
+					tblScan.HandleCols = NewIntHandleCols(col)
+					clonedTblScan.Schema().Append(col)
+					clonedTblScan.(*PhysicalTableScan).Columns = append(clonedTblScan.(*PhysicalTableScan).Columns, pk)
 				} else if tblInfo.IsCommonHandle {
 					idxInfo := tblInfo.GetPrimaryKey()
 					tblScan.HandleCols = NewCommonHandleCols(p.SCtx().GetSessionVars().StmtCtx, tblInfo, idxInfo, tblScan.tblCols)
 					for i := 0; i < tblScan.HandleCols.NumCols(); i++ {
-						clonedTblScan.(*PhysicalTableScan).schema.Append(tblScan.HandleCols.GetCol(i))
+						clonedTblScan.Schema().Append(tblScan.HandleCols.GetCol(i))
 						clonedTblScan.(*PhysicalTableScan).Columns = append(clonedTblScan.(*PhysicalTableScan).Columns, tblScan.HandleCols.GetCol(i).ToInfo())
 					}
 				} else {
-					clonedTblScan.(*PhysicalTableScan).schema.Append(tblScan.HandleCols.GetCol(0))
+					clonedTblScan.Schema().Append(tblScan.HandleCols.GetCol(0))
 					clonedTblScan.(*PhysicalTableScan).Columns = append(clonedTblScan.(*PhysicalTableScan).Columns, model.NewExtraHandleColInfo())
 				}
 			}