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