planner: fix unexpected behavior of UPDATE (pingcap#12597)

winoros · Nov 4, 2019 · 8c9b3d5 · 8c9b3d5
1 parent ae6b36c
commit 8c9b3d5
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git a/executor/update_test.go b/executor/update_test.go
@@ -216,3 +216,14 @@ func (s *testUpdateSuite) TestUpdateWithAutoidSchema(c *C) {
 		tk.MustQuery(tt.query).Check(tt.result)
 	}
 }
+
+func (s *testUpdateSuite) TestUpdateWithSubquery(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t1(id varchar(30) not null, status varchar(1) not null default 'N', id2 varchar(30))")
+	tk.MustExec("create table t2(id varchar(30) not null, field varchar(4) not null)")
+	tk.MustExec("insert into t1 values('abc', 'F', 'abc')")
+	tk.MustExec("insert into t2 values('abc', 'MAIN')")
+	tk.MustExec("update t1 set status = 'N' where status = 'F' and (id in (select id from t2 where field = 'MAIN') or id2 in (select id from t2 where field = 'main'))")
+	tk.MustQuery("select * from t1").Check(testkit.Rows("abc N abc"))
+}
diff --git a/planner/core/logical_plan_builder.go b/planner/core/logical_plan_builder.go
@@ -2263,15 +2263,16 @@ func (b *planBuilder) buildUpdate(update *ast.UpdateStmt) (Plan, error) {
 		b.visitInfo = appendVisitInfo(b.visitInfo, mysql.SelectPriv, dbName, t.Name.L, "")
 	}
 
-	oldSchemaLen := p.Schema().Len()
+	oldSchema := p.Schema().Clone()
 	if sel.Where != nil {
 		p, err = b.buildSelection(p, sel.Where, nil)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
 	}
-	if update.Order != nil {
-		p, err = b.buildSort(p, update.Order.Items, nil)
+
+	if sel.OrderBy != nil {
+		p, err = b.buildSort(p, sel.OrderBy.Items, nil)
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
@@ -2283,10 +2284,9 @@ func (b *planBuilder) buildUpdate(update *ast.UpdateStmt) (Plan, error) {
 		}
 	}
 	// TODO: expression rewriter should not change the output columns. We should cut the columns here.
-	if p.Schema().Len() != oldSchemaLen {
-		proj := LogicalProjection{Exprs: expression.Column2Exprs(p.Schema().Columns[:oldSchemaLen])}.init(b.ctx)
-		proj.SetSchema(expression.NewSchema(make([]*expression.Column, oldSchemaLen)...))
-		copy(proj.schema.Columns, p.Schema().Columns[:oldSchemaLen])
+	if p.Schema().Len() != oldSchema.Len() {
+		proj := LogicalProjection{Exprs: expression.Column2Exprs(oldSchema.Columns)}.init(b.ctx)
+		proj.SetSchema(oldSchema)
 		proj.SetChildren(p)
 		p = proj
 	}

diff --git a/planner/core/logical_plan_test.go b/planner/core/logical_plan_test.go
@@ -957,7 +957,7 @@ func (s *testPlanSuite) TestPlanBuilder(c *C) {
 		},
 		{
 			sql:  "update t set a = 2 where b in (select c from t)",
-			plan: "LeftHashJoin{TableReader(Table(t))->IndexReader(Index(t.c_d_e)[[NULL,+inf]])->HashAgg}(Column#2,Column#25)->Projection->Update",
+			plan: "LeftHashJoin{TableReader(Table(t))->TableReader(Table(t))}(test.t.b,test.t.c)->Update",
 		},
 	}
 	for _, ca := range tests {