tiny refactor

Signed-off-by: crazycs520 <crazycs520@gmail.com>

executor/update.go

@@ -245,13 +245,7 @@ func (e *UpdateExec) Next(ctx context.Context, req *chunk.Chunk) error {
 
 func (e *UpdateExec) updateRows(ctx context.Context) (int, error) {
 	fields := retTypes(e.children[0])
-	colsInfo := make([]*table.Column, len(fields))
-	for _, content := range e.tblColPosInfos {
-		tbl := e.tblID2table[content.TblID]
-		for i, c := range tbl.WritableCols() {
-			colsInfo[content.Start+i] = c
-		}
-	}
+	colsInfo := plannercore.GetUpdateColumnsInfo(e.tblID2table, e.tblColPosInfos, len(fields))
 	globalRowIdx := 0
 	chk := newFirstChunk(e.children[0])
 	if !e.allAssignmentsAreConstant {
@@ -543,9 +537,9 @@ func (e *updateRuntimeStats) Tp() int {
 
 // GetFKChecks implements WithForeignKeyTrigger interface.
 func (e *UpdateExec) GetFKChecks() []*FKCheckExec {
-	fkChecks := []*FKCheckExec{}
-	for _, fkcs := range e.fkChecks {
-		fkChecks = append(fkChecks, fkcs...)
+	fkChecks := make([]*FKCheckExec, 0, len(e.fkChecks))
+	for _, fkc := range e.fkChecks {
+		fkChecks = append(fkChecks, fkc...)
 	}
 	return fkChecks
 }

planner/core/foreign_key.go

@@ -139,21 +139,15 @@ func buildOnUpdateReferredFKChecks(is infoschema.InfoSchema, dbName string, tblI
 }
 
 func (updt *Update) buildTbl2UpdateColumns() map[int64]map[string]struct{} {
-	colsInfo := make([]*model.ColumnInfo, len(updt.SelectPlan.Schema().Columns))
-	for _, content := range updt.TblColPosInfos {
-		tbl := updt.tblID2Table[content.TblID]
-		for i, c := range tbl.WritableCols() {
-			colsInfo[content.Start+i] = c.ColumnInfo
-		}
-	}
+	colsInfo := GetUpdateColumnsInfo(updt.tblID2Table, updt.TblColPosInfos, len(updt.SelectPlan.Schema().Columns))
 	tblID2UpdateColumns := make(map[int64]map[string]struct{})
-	for tid := range updt.tblID2Table {
-		tblID2UpdateColumns[tid] = make(map[string]struct{})
-	}
 	for _, assign := range updt.OrderedList {
 		col := colsInfo[assign.Col.Index]
 		for _, content := range updt.TblColPosInfos {
 			if assign.Col.Index >= content.Start && assign.Col.Index < content.End {
+				if _, ok := tblID2UpdateColumns[content.TblID]; !ok {
+					tblID2UpdateColumns[content.TblID] = make(map[string]struct{})
+				}
 				tblID2UpdateColumns[content.TblID][col.Name.L] = struct{}{}
 				break
 			}

planner/core/logical_plan_builder.go

@@ -5385,6 +5385,18 @@ func (b *PlanBuilder) buildUpdate(ctx context.Context, update *ast.UpdateStmt) (
 	return updt, err
 }
 
+// GetUpdateColumnsInfo get the update columns info.
+func GetUpdateColumnsInfo(tblID2Table map[int64]table.Table, tblColPosInfos TblColPosInfoSlice, size int) []*table.Column {
+	colsInfo := make([]*table.Column, size)
+	for _, content := range tblColPosInfos {
+		tbl := tblID2Table[content.TblID]
+		for i, c := range tbl.WritableCols() {
+			colsInfo[content.Start+i] = c
+		}
+	}
+	return colsInfo
+}
+
 type tblUpdateInfo struct {
 	name                string
 	pkUpdated           bool