planner: UPDATE's select plan's output col IDs should be stable (#53268)

expression/BUILD.bazel

@@ -99,6 +99,7 @@ go_library(
         "//util/encrypt",
         "//util/generatedexpr",
         "//util/hack",
+        "//util/intset",
         "//util/logutil",
         "//util/mathutil",
         "//util/mock",
@@ -120,7 +121,6 @@ go_library(
         "@com_github_pkg_errors//:errors",
         "@com_github_tikv_client_go_v2//oracle",
         "@org_golang_x_exp//slices",
-        "@org_golang_x_tools//container/intsets",
         "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],

expression/util.go

@@ -36,10 +36,10 @@ import (
 	driver "github.com/pingcap/tidb/types/parser_driver"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/collate"
+	"github.com/pingcap/tidb/util/intset"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/sqlexec"
 	"go.uber.org/zap"
-	"golang.org/x/tools/container/intsets"
 )
 
 // cowExprRef is a copy-on-write slice ref util using in `ColumnSubstitute`
@@ -373,15 +373,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)
 	}
 	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))

planner/core/issuetest/BUILD.bazel

@@ -6,6 +6,6 @@ go_test(
     srcs = ["planner_issue_test.go"],
     flaky = True,
     race = "on",
-    shard_count = 15,
+    shard_count = 16,
     deps = ["//testkit"],
 )

planner/core/issuetest/planner_issue_test.go

@@ -244,21 +244,17 @@ func TestIssue50614(t *testing.T) {
 		testkit.Rows(
 			"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#7, Column#8, Column#9",
-			"      │ │ └─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#11, Column#12, Column#13",
-			"      │ │ └─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#15, Column#16, Column#17",
-			"      │   └─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",
 		),
 	)
 }
@@ -340,6 +336,36 @@ WHERE
 	tk.MustQuery(`show warnings`).Check(testkit.Rows())
 }
 
+// https://github.com/pingcap/tidb/issues/53236
+func TestIssue53236(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test;")
+	tk.MustExec("create table t1(id int primary key, a varchar(128));")
+	tk.MustExec("create table t2(id int primary key, b varchar(128), c varchar(128));")
+	tk.MustExec(`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;`)
+}
+
 func TestIssue52687(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)

planner/core/logical_plan_builder.go

@@ -6085,8 +6085,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

planner/core/rule_eliminate_projection.go

@@ -72,13 +72,6 @@ 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]) {
@@ -136,6 +129,11 @@ func doPhysicalProjectionElimination(p PhysicalPlan) 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
 }
 

sessionctx/stmtctx/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//parser/terror",
         "//util/disk",
         "//util/execdetails",
+        "//util/intset",
         "//util/memory",
         "//util/resourcegrouptag",
         "//util/topsql/stmtstats",

sessionctx/stmtctx/stmtctx.go

@@ -35,6 +35,7 @@ import (
 	"github.com/pingcap/tidb/parser/terror"
 	"github.com/pingcap/tidb/util/disk"
 	"github.com/pingcap/tidb/util/execdetails"
+	"github.com/pingcap/tidb/util/intset"
 	"github.com/pingcap/tidb/util/memory"
 	"github.com/pingcap/tidb/util/resourcegrouptag"
 	"github.com/pingcap/tidb/util/topsql/stmtstats"
@@ -374,7 +375,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
 
 	// RangeFallback indicates that building complete ranges exceeds the memory limit so it falls back to less accurate ranges such as full range.
 	RangeFallback bool

util/intset/BUILD.bazel

@@ -0,0 +1,26 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "intset",
+    srcs = ["fast_int_set.go"],
+    importpath = "github.com/pingcap/tidb/util/intset",
+    visibility = ["//visibility:public"],
+    deps = ["@org_golang_x_tools//container/intsets"],
+)
+
+go_test(
+    name = "intset_test",
+    timeout = "short",
+    srcs = [
+        "fast_int_set_bench_test.go",
+        "fast_int_set_test.go",
+    ],
+    embed = [":intset"],
+    flaky = True,
+    deps = [
+        "@com_github_stretchr_testify//require",
+        "@org_golang_x_exp//maps",
+        "@org_golang_x_exp//slices",
+        "@org_golang_x_tools//container/intsets",
+    ],
+)