planner: UPDATE's select plan's output col IDs should be stable (ping…

…cap#53268)

close pingcap#53236

build/BUILD.bazel

@@ -178,7 +178,6 @@ nogo(
            }) +
            select({
                "//build:without_rbe": [
-                   "//build/linter/filepermission",
                ],
                "//conditions:default": [],
            }),

pkg/expression/BUILD.bazel

@@ -124,7 +124,6 @@ go_library(
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_tipb//go-tipb",
         "@com_github_tikv_client_go_v2//oracle",
-        "@org_golang_x_tools//container/intsets",
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],

pkg/expression/util.go

@@ -36,9 +36,9 @@ import (
 	driver "github.com/pingcap/tidb/pkg/types/parser_driver"
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/collate"
+	"github.com/pingcap/tidb/pkg/util/intset"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"go.uber.org/zap"
-	"golang.org/x/tools/container/intsets"
 )
 
 // cowExprRef is a copy-on-write slice ref util using in `ColumnSubstitute`
@@ -372,15 +372,15 @@ func ExtractColumnsAndCorColumnsFromExpressions(result []*Column, list []Express
 }
 
 // ExtractColumnSet extracts the different values of `UniqueId` for columns in expressions.
-func ExtractColumnSet(exprs ...Expression) *intsets.Sparse {
-	set := &intsets.Sparse{}
+func ExtractColumnSet(exprs ...Expression) intset.FastIntSet {
+	set := intset.NewFastIntSet()
 	for _, expr := range exprs {
-		extractColumnSet(expr, set)
+		extractColumnSet(expr, &set)
 	}
 	return set
 }
 
-func extractColumnSet(expr Expression, set *intsets.Sparse) {
+func extractColumnSet(expr Expression, set *intset.FastIntSet) {
 	switch v := expr.(type) {
 	case *Column:
 		set.Insert(int(v.UniqueID))

pkg/planner/core/logical_plan_builder.go

@@ -6090,8 +6090,8 @@ 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)
+		if cols := expression.ExtractColumnSet(newExpr); cols.Len() > 0 {
+			b.ctx.GetSessionVars().StmtCtx.ColRefFromUpdatePlan.UnionWith(cols)
 		}
 		newList = append(newList, &expression.Assignment{Col: col, ColName: name.ColName, Expr: newExpr})
 		dbName := name.DBName.L

pkg/planner/core/rule_eliminate_projection.go

@@ -82,13 +82,6 @@ func canProjectionBeEliminatedStrict(p *PhysicalProjection) bool {
 	if p.Schema().Len() != child.Schema().Len() {
 		return false
 	}
-	for _, ref := range p.SCtx().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.EqualColumn(child.Schema().Columns[i]) {
@@ -146,6 +139,11 @@ func doPhysicalProjectionElimination(p base.PhysicalPlan) base.PhysicalPlan {
 			childProj.SetSchema(p.Schema())
 		}
 	}
+	for i, col := range p.Schema().Columns {
+		if p.SCtx().GetSessionVars().StmtCtx.ColRefFromUpdatePlan.Has(int(col.UniqueID)) && !child.Schema().Columns[i].Equal(nil, col) {
+			return p
+		}
+	}
 	return child
 }
 

pkg/sessionctx/stmtctx/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
         "//pkg/util/execdetails",
         "//pkg/util/hint",
         "//pkg/util/intest",
+        "//pkg/util/intset",
         "//pkg/util/linter/constructor",
         "//pkg/util/memory",
         "//pkg/util/nocopy",

pkg/sessionctx/stmtctx/stmtctx.go

@@ -41,6 +41,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/execdetails"
 	"github.com/pingcap/tidb/pkg/util/hint"
 	"github.com/pingcap/tidb/pkg/util/intest"
+	"github.com/pingcap/tidb/pkg/util/intset"
 	"github.com/pingcap/tidb/pkg/util/linter/constructor"
 	"github.com/pingcap/tidb/pkg/util/memory"
 	"github.com/pingcap/tidb/pkg/util/nocopy"
@@ -362,7 +363,7 @@ type StatementContext struct {
 	// UseDynamicPruneMode indicates whether use UseDynamicPruneMode in query stmt
 	UseDynamicPruneMode bool
 	// ColRefFromPlan mark the column ref used by assignment in update statement.
-	ColRefFromUpdatePlan []int64
+	ColRefFromUpdatePlan intset.FastIntSet
 
 	// IsExplainAnalyzeDML is true if the statement is "explain analyze DML executors", before responding the explain
 	// results to the client, the transaction should be committed first. See issue #37373 for more details.

tests/integrationtest/r/planner/core/casetest/integration.result

@@ -1803,18 +1803,14 @@ explain format = brief update tt, (select 1 as c1 ,2 as c2 ,3 as c3, 4 as c4 uni
 id	estRows	task	access object	operator info
 Update	N/A	root		N/A
 └─Projection	0.00	root		test.tt.a, test.tt.b, test.tt.c, test.tt.d, test.tt.e, Column#18, Column#19, Column#20, Column#21
-  └─Projection	0.00	root		test.tt.a, test.tt.b, test.tt.c, test.tt.d, test.tt.e, Column#18, Column#19, Column#20, Column#21
-    └─IndexJoin	0.00	root		inner join, inner:TableReader, outer key:Column#20, Column#21, inner key:test.tt.c, test.tt.d, equal cond:eq(Column#20, test.tt.c), eq(Column#21, test.tt.d), other cond:or(or(and(eq(Column#20, 11), eq(test.tt.d, 111)), and(eq(Column#20, 22), eq(test.tt.d, 222))), or(and(eq(Column#20, 33), eq(test.tt.d, 333)), and(eq(Column#20, 44), eq(test.tt.d, 444)))), or(or(and(eq(test.tt.c, 11), eq(Column#21, 111)), and(eq(test.tt.c, 22), eq(Column#21, 222))), or(and(eq(test.tt.c, 33), eq(Column#21, 333)), and(eq(test.tt.c, 44), eq(Column#21, 444))))
-      ├─Union(Build)	0.00	root		
-      │ ├─Projection	0.00	root		Column#6->Column#18, Column#7->Column#19, Column#8->Column#20, Column#9->Column#21
-      │ │ └─Projection	0.00	root		1->Column#6, 2->Column#7, 3->Column#8, 4->Column#9
-      │ │   └─TableDual	0.00	root		rows:0
-      │ ├─Projection	0.00	root		Column#10->Column#18, Column#11->Column#19, Column#12->Column#20, Column#13->Column#21
-      │ │ └─Projection	0.00	root		2->Column#10, 3->Column#11, 4->Column#12, 5->Column#13
-      │ │   └─TableDual	0.00	root		rows:0
-      │ └─Projection	0.00	root		Column#14->Column#18, Column#15->Column#19, Column#16->Column#20, Column#17->Column#21
-      │   └─Projection	0.00	root		3->Column#14, 4->Column#15, 5->Column#16, 6->Column#17
-      │     └─TableDual	0.00	root		rows:0
-      └─TableReader(Probe)	0.00	root		data:Selection
-        └─Selection	0.00	cop[tikv]		or(or(and(eq(test.tt.c, 11), eq(test.tt.d, 111)), and(eq(test.tt.c, 22), eq(test.tt.d, 222))), or(and(eq(test.tt.c, 33), eq(test.tt.d, 333)), and(eq(test.tt.c, 44), eq(test.tt.d, 444)))), or(or(eq(test.tt.c, 11), eq(test.tt.c, 22)), or(eq(test.tt.c, 33), eq(test.tt.c, 44))), or(or(eq(test.tt.d, 111), eq(test.tt.d, 222)), or(eq(test.tt.d, 333), eq(test.tt.d, 444)))
-          └─TableRangeScan	0.00	cop[tikv]	table:tt	range: decided by [eq(test.tt.c, Column#20) eq(test.tt.d, Column#21)], keep order:false, stats:pseudo
+  └─IndexJoin	0.00	root		inner join, inner:TableReader, outer key:Column#20, Column#21, inner key:test.tt.c, test.tt.d, equal cond:eq(Column#20, test.tt.c), eq(Column#21, test.tt.d), other cond:or(or(and(eq(Column#20, 11), eq(test.tt.d, 111)), and(eq(Column#20, 22), eq(test.tt.d, 222))), or(and(eq(Column#20, 33), eq(test.tt.d, 333)), and(eq(Column#20, 44), eq(test.tt.d, 444)))), or(or(and(eq(test.tt.c, 11), eq(Column#21, 111)), and(eq(test.tt.c, 22), eq(Column#21, 222))), or(and(eq(test.tt.c, 33), eq(Column#21, 333)), and(eq(test.tt.c, 44), eq(Column#21, 444))))
+    ├─Union(Build)	0.00	root		
+    │ ├─Projection	0.00	root		1->Column#18, 2->Column#19, 3->Column#20, 4->Column#21
+    │ │ └─TableDual	0.00	root		rows:0
+    │ ├─Projection	0.00	root		2->Column#18, 3->Column#19, 4->Column#20, 5->Column#21
+    │ │ └─TableDual	0.00	root		rows:0
+    │ └─Projection	0.00	root		3->Column#18, 4->Column#19, 5->Column#20, 6->Column#21
+    │   └─TableDual	0.00	root		rows:0
+    └─TableReader(Probe)	0.00	root		data:Selection
+      └─Selection	0.00	cop[tikv]		or(or(and(eq(test.tt.c, 11), eq(test.tt.d, 111)), and(eq(test.tt.c, 22), eq(test.tt.d, 222))), or(and(eq(test.tt.c, 33), eq(test.tt.d, 333)), and(eq(test.tt.c, 44), eq(test.tt.d, 444)))), or(or(eq(test.tt.c, 11), eq(test.tt.c, 22)), or(eq(test.tt.c, 33), eq(test.tt.c, 44))), or(or(eq(test.tt.d, 111), eq(test.tt.d, 222)), or(eq(test.tt.d, 333), eq(test.tt.d, 444)))
+        └─TableRangeScan	0.00	cop[tikv]	table:tt	range: decided by [eq(test.tt.c, Column#20) eq(test.tt.d, Column#21)], keep order:false, stats:pseudo

tests/integrationtest/r/planner/core/cbo.result

@@ -3,14 +3,13 @@ create table t(a int, b int);
 explain update t t1, (select distinct b from t) t2 set t1.b = t2.b;
 id	estRows	task	access object	operator info
 Update_7	N/A	root		N/A
-└─Projection_9	80000000.00	root		planner__core__cbo.t.a, planner__core__cbo.t.b, planner__core__cbo.t._tidb_rowid, planner__core__cbo.t.b
-  └─HashJoin_10	80000000.00	root		CARTESIAN inner join
-    ├─HashAgg_18(Build)	8000.00	root		group by:planner__core__cbo.t.b, funcs:firstrow(planner__core__cbo.t.b)->planner__core__cbo.t.b
-    │ └─TableReader_19	8000.00	root		data:HashAgg_14
-    │   └─HashAgg_14	8000.00	cop[tikv]		group by:planner__core__cbo.t.b, 
-    │     └─TableFullScan_17	10000.00	cop[tikv]	table:t	keep order:false, stats:pseudo
-    └─TableReader_13(Probe)	10000.00	root		data:TableFullScan_12
-      └─TableFullScan_12	10000.00	cop[tikv]	table:t1	keep order:false, stats:pseudo
+└─HashJoin_10	80000000.00	root		CARTESIAN inner join
+  ├─HashAgg_18(Build)	8000.00	root		group by:planner__core__cbo.t.b, funcs:firstrow(planner__core__cbo.t.b)->planner__core__cbo.t.b
+  │ └─TableReader_19	8000.00	root		data:HashAgg_14
+  │   └─HashAgg_14	8000.00	cop[tikv]		group by:planner__core__cbo.t.b, 
+  │     └─TableFullScan_17	10000.00	cop[tikv]	table:t	keep order:false, stats:pseudo
+  └─TableReader_13(Probe)	10000.00	root		data:TableFullScan_12
+    └─TableFullScan_12	10000.00	cop[tikv]	table:t1	keep order:false, stats:pseudo
 drop table if exists tb1, tb2;
 create table tb1(a int, b int, primary key(a));
 create table tb2 (a int, b int, c int, d datetime, primary key(c),key idx_u(a));

tests/integrationtest/r/planner/core/issuetest/planner_issue.result

@@ -520,6 +520,29 @@ t_q1 as ref_14
 where (ref_14.c_z like 'o%fiah')))
 where (t_kg74.c_obnq8s7_s2 = case when (t_kg74.c_a1tv2 is NULL) then t_kg74.c_g else t_kg74.c_obnq8s7_s2 end
 );
+drop table if exists t1, t2;
+create table t1(id int primary key, a varchar(128));
+create table t2(id int primary key, b varchar(128), c varchar(128));
+UPDATE
+t1
+SET
+t1.a = IFNULL(
+(
+SELECT
+t2.c
+FROM
+t2
+WHERE
+t2.b = t1.a
+ORDER BY
+t2.b DESC,
+t2.c DESC
+LIMIT
+1
+), ''
+)
+WHERE
+t1.id = 1;
 drop table if exists t0, t1;
 CREATE TABLE t0(c0 NUMERIC);
 CREATE TABLE t1(c0 NUMERIC);

tests/integrationtest/r/planner/core/rule_constant_propagation.result

@@ -149,15 +149,13 @@ create table s (id int, name varchar(10));
 explain Update t, (select * from s where s.id>1) tmp set t.name=tmp.name where t.id=tmp.id;
 id	estRows	task	access object	operator info
 Update_8	N/A	root		N/A
-└─Projection_11	4166.67	root		planner__core__rule_constant_propagation.t.id, planner__core__rule_constant_propagation.t.name, planner__core__rule_constant_propagation.t._tidb_rowid, planner__core__rule_constant_propagation.s.id, planner__core__rule_constant_propagation.s.name
-  └─HashJoin_12	4166.67	root		inner join, equal:[eq(planner__core__rule_constant_propagation.t.id, planner__core__rule_constant_propagation.s.id)]
-    ├─Projection_17(Build)	3333.33	root		planner__core__rule_constant_propagation.s.id, planner__core__rule_constant_propagation.s.name
-    │ └─TableReader_20	3333.33	root		data:Selection_19
-    │   └─Selection_19	3333.33	cop[tikv]		gt(planner__core__rule_constant_propagation.s.id, 1), not(isnull(planner__core__rule_constant_propagation.s.id))
-    │     └─TableFullScan_18	10000.00	cop[tikv]	table:s	keep order:false, stats:pseudo
-    └─TableReader_16(Probe)	3333.33	root		data:Selection_15
-      └─Selection_15	3333.33	cop[tikv]		gt(planner__core__rule_constant_propagation.t.id, 1), not(isnull(planner__core__rule_constant_propagation.t.id))
-        └─TableFullScan_14	10000.00	cop[tikv]	table:t	keep order:false, stats:pseudo
+└─HashJoin_12	4166.67	root		inner join, equal:[eq(planner__core__rule_constant_propagation.t.id, planner__core__rule_constant_propagation.s.id)]
+  ├─TableReader_20(Build)	3333.33	root		data:Selection_19
+  │ └─Selection_19	3333.33	cop[tikv]		gt(planner__core__rule_constant_propagation.s.id, 1), not(isnull(planner__core__rule_constant_propagation.s.id))
+  │   └─TableFullScan_18	10000.00	cop[tikv]	table:s	keep order:false, stats:pseudo
+  └─TableReader_16(Probe)	3333.33	root		data:Selection_15
+    └─Selection_15	3333.33	cop[tikv]		gt(planner__core__rule_constant_propagation.t.id, 1), not(isnull(planner__core__rule_constant_propagation.t.id))
+      └─TableFullScan_14	10000.00	cop[tikv]	table:t	keep order:false, stats:pseudo
 drop table if exists t, s;
 create table t (id int, name varchar(10));
 create table s (id int, name varchar(10));

tests/integrationtest/t/planner/core/issuetest/planner_issue.test

@@ -379,6 +379,31 @@ update t_kg74 set
 where (t_kg74.c_obnq8s7_s2 = case when (t_kg74.c_a1tv2 is NULL) then t_kg74.c_g else t_kg74.c_obnq8s7_s2 end
       );
 
+# https://github.com/pingcap/tidb/issues/53236
+drop table if exists t1, t2;
+create table t1(id int primary key, a varchar(128));
+create table t2(id int primary key, b varchar(128), c varchar(128));
+UPDATE
+    t1
+SET
+    t1.a = IFNULL(
+            (
+                SELECT
+                    t2.c
+                FROM
+                    t2
+                WHERE
+                    t2.b = t1.a
+                ORDER BY
+                    t2.b DESC,
+                    t2.c DESC
+                LIMIT
+                    1
+            ), ''
+        )
+WHERE
+    t1.id = 1;
+
 # https://github.com/pingcap/tidb/issues/49109
 drop table if exists t0, t1;
 CREATE TABLE t0(c0 NUMERIC);