plan: fix insert onduplicate key update

executor/insert_test.go

@@ -81,6 +81,20 @@ func (s *testSuite) TestInsertOnDuplicateKey(c *C) {
 	tk.MustQuery(`select * from t1`).Check(testkit.Rows("1 400"))
 	tk.MustExec(`insert into t1 select * from t2 on duplicate key update b1 = values(b1) + b2;`)
 	tk.MustQuery(`select * from t1`).Check(testkit.Rows("1 400"))
+
+	tk.MustExec(`drop table if exists t;`)
+	tk.MustExec(`create table t(k1 bigint, k2 bigint, val bigint, primary key(k1, k2));`)
+	tk.MustExec(`insert into t (val, k1, k2) values (3, 1, 2);`)
+	tk.MustQuery(`select * from t;`).Check(testkit.Rows(`1 2 3`))
+	tk.MustExec(`insert into t (val, k1, k2) select c, a, b from (select 1 as a, 2 as b, 4 as c) tmp on duplicate key update val = tmp.c;`)
+	tk.MustQuery(`select * from t;`).Check(testkit.Rows(`1 2 4`))
+
+	tk.MustExec(`drop table if exists t;`)
+	tk.MustExec(`create table t(k1 double, k2 double, v double, primary key(k1, k2));`)
+	tk.MustExec(`insert into t (v, k1, k2) select c, a, b from (select "3" c, "1" a, "2" b) tmp on duplicate key update v=c;`)
+	tk.MustQuery(`select * from t;`).Check(testkit.Rows(`1 2 3`))
+	tk.MustExec(`insert into t (v, k1, k2) select c, a, b from (select "3" c, "1" a, "2" b) tmp on duplicate key update v=c;`)
+	tk.MustQuery(`select * from t;`).Check(testkit.Rows(`1 2 3`))
 }
 
 func (s *testSuite) TestInsertWrongValueForField(c *C) {

plan/planbuilder.go

@@ -1187,7 +1187,40 @@ func (b *planBuilder) buildSelectPlanOfInsert(insert *ast.InsertStmt, insertPlan
 		return errors.Trace(err)
 	}
 
-	insertPlan.Schema4OnDuplicate = expression.MergeSchema(insertPlan.tableSchema, insertPlan.SelectPlan.Schema())
+	tableCols := insertPlan.Table.Cols()
+	var insertCols []*table.Column
+
+	if len(insertPlan.Columns) > 0 {
+		colNames := make([]string, 0, len(tableCols))
+		for _, col := range insertPlan.Columns {
+			colNames = append(colNames, col.Name.O)
+		}
+		for _, col := range insertPlan.GenCols.Columns {
+			colNames = append(colNames, col.Name.O)
+		}
+		insertCols, err = table.FindCols(tableCols, colNames, insertPlan.Table.Meta().PKIsHandle)
+		if err != nil {
+			return errors.Trace(err)
+		}
+	} else {
+		insertCols = tableCols
+	}
+
+	schema4NewRow := expression.NewSchema(make([]*expression.Column, len(tableCols))...)
+	for i, selCol := range insertPlan.SelectPlan.Schema().Columns {
+		ordinal := insertCols[i].Offset
+		schema4NewRow.Columns[ordinal] = &expression.Column{}
+		*schema4NewRow.Columns[ordinal] = *selCol
+
+		schema4NewRow.Columns[ordinal].RetType = &types.FieldType{}
+		*schema4NewRow.Columns[ordinal].RetType = insertCols[i].FieldType
+	}
+	for i := range schema4NewRow.Columns {
+		if schema4NewRow.Columns[i] == nil {
+			schema4NewRow.Columns[i] = &expression.Column{UniqueID: insertPlan.ctx.GetSessionVars().AllocPlanColumnID()}
+		}
+	}
+	insertPlan.Schema4OnDuplicate = expression.MergeSchema(insertPlan.tableSchema, schema4NewRow)
 	return nil
 }