Repository: pingcap/tidb (Go)

Issue:
Can't find column in update statement and wrong update results
## Bug Report

Please answer these questions before submitting your issue. Thanks!

### 1. Minimal reproduce step (Required)

```sql
drop table if exists tb1, tb2;
create table tb1(user_no int ,balance int,primary key(user_no));
create table tb2 (user_no int ,balance int,seq_no int,dt datetime,primary key(seq_no),key idx_u(user_no));
insert into tb1 values(1,100),(2,100),(3,200),(4,1000),(5,232),(6,111);
insert into tb2 values(1,99,1,now()),(2,101,2,now()),(3,800,3,now()),(4,6000,4,now()),(5,232,5,now()),(6,11132,6,now());
insert into tb2 values(1,999,7,now()),(2,1013,8,now()),(3,83200,9,now()),(4,632000,10,now()),(5,2332,11,now()),(6,121132,12,now());
update tb1 set tb1.balance=(select tb2.balance from tb2 where tb2.user_no=tb1.user_no order by seq_no desc limit 1);
update tb1 inner join tb2 on tb1.user_no = tb2.user_no set tb1.balance = tb2.balance order by tb2.seq_no desc;
select * from tb1;
```
<!-- a step by step guide for reproducing the bug. -->

### 2. What did you expect to see? (Required)

The two update statements should be successful.

### 3. What did you see instead (Required)

The first one return an error, "ERROR 1105 (HY000): Can't find column test.tb2.balance in schema Column: [test.tb1.user_no,test.tb1.balance,test.tb2.balance] Unique key: []"

If I remove the first update statement, the second update statement will succeed, but it updates the wrong values.

```
mysql> select * from tb1;
+---------+---------+
| user_no | balance |
+---------+---------+
|       1 |     999 |
|       2 |    1013 |
|       3 |   83200 |
|       4 |  632000 |
|       5 |     232 |
|       6 |  121132 |
+---------+---------+
6 rows in set (0.00 sec)
```

The balance of user_no = 5 should be 2332.

If I run the second update statement again, I get

```
mysql> select * from tb1;
+---------+---------+
| user_no | balance |
+---------+---------+
|       1 |      99 |
|       2 |     101 |
|       3 |     800 |
|       4 |    6000 |
|       5 |    2332 |
|       6 |   11132 |
+---------+---------+
6 rows in set (0.00 sec)
```

It seems the order by seq_no clause does not work.

I try the first statement in MySQL 8.0 it gets the correct results, and the second one returns an error, "ERROR 1221 (HY000): Incorrect usage of UPDATE and ORDER BY".

### 4. What is your TiDB version? (Required)

v6.1.0
<!-- Paste the output of SELECT tidb_version() -->




Discussion:
For the first update error, it seems there's something wrong in the logical plan builder:

```
github.com/pingcap/tidb/expression.(*Column).resolveIndices
\t/Users/yusp/work/tidb/expression/column.go:503
github.com/pingcap/tidb/expression.(*Column).ResolveIndices
\t/Users/yusp/work/tidb/expression/column.go:496
github.com/pingcap/tidb/planner/core.(*Update).ResolveIndices
\t/Users/yusp/work/tidb/planner/core/resolve_indices.go:586
github.com/pingcap/tidb/planner/core.(*PlanBuilder).buildUpdate
\t/Users/yusp/work/tidb/planner/core/logical_plan_builder.go:5145
github.com/pingcap/tidb/planner/core.(*PlanBuilder).Build
\t/Users/yusp/work/tidb/planner/core/planbuilder.go:705
github.com/pingcap/tidb/planner.optimize
\t/Users/yusp/work/tidb/planner/optimize.go:354
github.com/pingcap/tidb/planner.Optimize
\t/Users/yusp/work/tidb/planner/optimize.go:227
github.com/pingcap/tidb/planner/core.(*PlanBuilder).buildExplain
\t/Users/yusp/work/tidb/planner/core/planbuilder.go:4526
github.com/pingcap/tidb/planner/core.(*PlanBuilder).Build
\t/Users/yusp/work/tidb/planner/core/planbuilder.go:680
github.com/pingcap/tidb/planner.optimize
\t/Users/yusp/work/tidb/planner/optimize.go:354
github.com/pingcap/tidb/planner.Optimize
\t/Users/yusp/work/tidb/planner/optimize.go:227
github.com/pingcap/tidb/executor.(*Compiler).Compile
\t/Users/yusp/work/tidb/executor/compiler.go:81
github.com/pingcap/tidb/session.(*session).ExecuteStmt
\t/Users/yusp/work/tidb/session/session.go:1916
github.com/pingcap/tidb/server.(*TiDBContext).ExecuteStmt
\t/Users/yusp/work/tidb/server/driver_tidb.go:230
github.com/pingcap/tidb/server.(*clientConn).handleStmt
\t/Users/yusp/work/tidb/server/conn.go:2022
github.com/pingcap/tidb/server.(*clientConn).handleQuery
\t/Users/yusp/work/tidb/server/conn.go:1876
github.com/pingcap/tidb/server.(*clientConn).dispatch
\t/Users/yusp/work/tidb/server/conn.go:1371
github.com/pingcap/tidb/server.(*clientConn).Run
\t/Users/yusp/work/tidb/server/conn.go:1121
github.com/pingcap/tidb/server.(*Server).onConn
\t/Users/yusp/work/tidb/server/server.go:559
runtime.goexit
\t/usr/local/go/src/runtime/asm_arm64.s:1165
```
Tried another update statement
```sql
update tb1,(select tb2.balance ,user_no from tb2  order by  seq_no desc limit 1 )  tb3 set tb1.balance=tb3.balance where tb1.user_no=tb3.user_no;
```
get a panic,

```
github.com/pingcap/tidb/server.(*clientConn).Run.func1
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/server/conn.go:1059
runtime.gopanic
\t/usr/local/go/src/runtime/panic.go:838
github.com/pingcap/tidb/executor.(*ExecStmt).Exec.func1
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/executor/adapter.go:371
runtime.gopanic
\t/usr/local/go/src/runtime/panic.go:838
runtime.goPanicIndex
\t/usr/local/go/src/runtime/panic.go:89
github.com/pingcap/tidb/util/chunk.(*Column).IsNull
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/util/chunk/column.go:170
github.com/pingcap/tidb/util/chunk.Row.IsNull
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/util/chunk/row.go:206
github.com/pingcap/tidb/util/chunk.Row.GetDatum
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/util/chunk/row.go:132
github.com/pingcap/tidb/util/chunk.Row.GetDatumRow
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/util/chunk/row.go:121
github.com/pingcap/tidb/executor.(*UpdateExec).updateRows
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/executor/update.go:284
github.com/pingcap/tidb/executor.(*UpdateExec).Next
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/executor/update.go:225
github.com/pingcap/tidb/executor.Next
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/executor/executor.go:319
github.com/pingcap/tidb/executor.(*ExecStmt).handleNoDelayExecutor
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/executor/adapter.go:665
github.com/pingcap/tidb/executor.(*ExecStmt).handleNoDelay
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/executor/adapter.go:516
github.com/pingcap/tidb/executor.(*ExecStmt).Exec
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/executor/adapter.go:465
github.com/pingcap/tidb/session.runStmt
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/session/session.go:2069
github.com/pingcap/tidb/session.(*session).ExecuteStmt
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/session/session.go:1938
github.com/pingcap/tidb/server.(*TiDBContext).ExecuteStmt
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/server/driver_tidb.go:230
github.com/pingcap/tidb/server.(*clientConn).handleStmt
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/server/conn.go:2022
github.com/pingcap/tidb/server.(*clientConn).handleQuery
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/server/conn.go:1876
github.com/pingcap/tidb/server.(*clientConn).dispatch
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/server/conn.go:1371
github.com/pingcap/tidb/server.(*clientConn).Run
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/server/conn.go:1121
github.com/pingcap/tidb/server.(*Server).onConn
\t/home/jenkins/agent/workspace/build-common/go/src/github.com/pingcap/tidb/server/server.go:559
```
But it doesn't appear in mocktikv, only with tikv cluster could meet this panic.
the reason is the same as: https://github.com/pingcap/tidb/pull/10962, we should avoid projection elimination in buildUpdate.

```
assignment {
  expr: depend on -> column-7 (comes from sub-query's output)
}

projection                                         apply join
column-7 <- column-4 --------------------- column-4 & ...

once the projection is eliminated here, the assignment couldn't resolve its self-expression anymore.
```



Patch:
diff --git a/planner/core/logical_plan_builder.go b/planner/core/logical_plan_builder.go
index bdc3922858b14..50b555db9ad04 100644
--- a/planner/core/logical_plan_builder.go
+++ b/planner/core/logical_plan_builder.go
@@ -5300,7 +5300,7 @@ func (b *PlanBuilder) buildUpdate(ctx context.Context, update *ast.UpdateStmt) (
 	}.Init(b.ctx)
 	updt.names = p.OutputNames()
 	// We cannot apply projection elimination when building the subplan, because
-	// columns in orderedList cannot be resolved.
+	// columns in orderedList cannot be resolved. (^flagEliminateProjection should also be applied in postOptimize)
 	updt.SelectPlan, _, err = DoOptimize(ctx, b.ctx, b.optFlag&^flagEliminateProjection, p)
 	if err != nil {
 		return nil, err
@@ -5549,6 +5549,9 @@ func (b *PlanBuilder) buildUpdateLists(ctx context.Context, tableList []*ast.Tab
 			allAssignmentsAreConstant = false
 		}
 		p = np
+		if col, ok := newExpr.(*expression.Column); ok {
+			b.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan = append(b.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan, col.UniqueID)
+		}
 		newList = append(newList, &expression.Assignment{Col: col, ColName: name.ColName, Expr: newExpr})
 		dbName := name.DBName.L
 		// To solve issue#10028, we need to get database name by the table alias name.
diff --git a/planner/core/optimizer.go b/planner/core/optimizer.go
index d43a9bfb28c21..75908649e1056 100644
--- a/planner/core/optimizer.go
+++ b/planner/core/optimizer.go
@@ -373,6 +373,8 @@ func mergeContinuousSelections(p PhysicalPlan) {
 }
 
 func postOptimize(sctx sessionctx.Context, plan PhysicalPlan) PhysicalPlan {
+	// some cases from update optimize will require avoiding projection elimination.
+	// see comments ahead of call of DoOptimize in function of buildUpdate().
 	plan = eliminatePhysicalProjection(plan)
 	plan = InjectExtraProjection(plan)
 	mergeContinuousSelections(plan)
diff --git a/planner/core/rule_eliminate_projection.go b/planner/core/rule_eliminate_projection.go
index 37d140f8cbd15..83936e7cd9e94 100644
--- a/planner/core/rule_eliminate_projection.go
+++ b/planner/core/rule_eliminate_projection.go
@@ -72,6 +72,13 @@ func canProjectionBeEliminatedStrict(p *PhysicalProjection) bool {
 	if p.Schema().Len() != child.Schema().Len() {
 		return false
 	}
+	for _, ref := range p.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan {
+		for _, one := range p.Schema().Columns {
+			if ref == one.UniqueID {
+				return false
+			}
+		}
+	}
 	for i, expr := range p.Exprs {
 		col, ok := expr.(*expression.Column)
 		if !ok || !col.Equal(nil, child.Schema().Columns[i]) {
diff --git a/sessionctx/stmtctx/stmtctx.go b/sessionctx/stmtctx/stmtctx.go
index f62df41de3615..7a290b264edc9 100644
--- a/sessionctx/stmtctx/stmtctx.go
+++ b/sessionctx/stmtctx/stmtctx.go
@@ -316,6 +316,8 @@ type StatementContext struct {
 	StatsLoadStatus map[model.TableItemID]string
 	// IsSyncStatsFailed indicates whether any failure happened during sync stats
 	IsSyncStatsFailed bool
+	// ColRefFromPlan mark the column ref used by assignment in update statement.
+	ColRefFromUpdatePlan []int64
 }
 
 // StmtHints are SessionVars related sql hints.


Test patch:
diff --git a/planner/core/cbo_test.go b/planner/core/cbo_test.go
index f1dce2c139163..d8d46c04d2aaa 100644
--- a/planner/core/cbo_test.go
+++ b/planner/core/cbo_test.go
@@ -726,6 +726,11 @@ func TestUpdateProjEliminate(t *testing.T) {
 	tk.MustExec("drop table if exists t")
 	tk.MustExec("create table t(a int, b int)")
 	tk.MustExec("explain update t t1, (select distinct b from t) t2 set t1.b = t2.b")
+
+	tk.MustExec("drop table if exists tb1, tb2")
+	tk.MustExec("create table tb1(a int, b int, primary key(a))")
+	tk.MustExec("create table tb2 (a int, b int, c int, d datetime, primary key(c),key idx_u(a));")
+	tk.MustExec("update tb1 set tb1.b=(select tb2.b from tb2 where tb2.a=tb1.a order by c desc limit 1);")
 }
 
 func TestTiFlashCostModel(t *testing.T) {