planner, expression: fix TableFullScan caused by extraHandleId not co…

…rrectly found (#45139) (#45141)

close #45036

expression/schema.go

@@ -18,6 +18,7 @@ import (
 	"strings"
 	"unsafe"
 
+	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/util/size"
 )
 
@@ -238,6 +239,17 @@ func (s *Schema) MemoryUsage() (sum int64) {
 	return
 }
 
+// GetExtraHandleColumn gets the extra handle column.
+func (s *Schema) GetExtraHandleColumn() *Column {
+	columnLen := len(s.Columns)
+	if columnLen > 0 && s.Columns[columnLen-1].ID == model.ExtraHandleID {
+		return s.Columns[columnLen-1]
+	} else if columnLen > 1 && s.Columns[columnLen-2].ID == model.ExtraHandleID {
+		return s.Columns[columnLen-2]
+	}
+	return nil
+}
+
 // MergeSchema will merge two schema into one schema. We shouldn't need to consider unique keys.
 // That will be processed in build_key_info.go.
 func MergeSchema(lSchema, rSchema *Schema) *Schema {

planner/core/issuetest/BUILD.bazel

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

planner/core/issuetest/planner_issue_test.go

@@ -46,6 +46,33 @@ func TestIssue42732(t *testing.T) {
 	tk.MustQuery("SELECT one.a, one.b as b2 FROM t1 one ORDER BY (SELECT two.b FROM t2 two WHERE two.a = one.b)").Check(testkit.Rows("1 1"))
 }
 
+// https://github.com/pingcap/tidb/issues/45036
+func TestIssue45036(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("CREATE TABLE ads_txn (\n" +
+		"  `cusno` varchar(10) NOT NULL,\n" +
+		"  `txn_dt` varchar(8) NOT NULL,\n" +
+		"  `unn_trno` decimal(22,0) NOT NULL,\n" +
+		"  `aml_cntpr_accno` varchar(64) DEFAULT NULL,\n" +
+		"  `acpayr_accno` varchar(35) DEFAULT NULL,\n" +
+		"  PRIMARY KEY (`cusno`,`txn_dt`,`unn_trno`) NONCLUSTERED\n" +
+		") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin\n" +
+		"PARTITION BY LIST COLUMNS(`txn_dt`)\n" +
+		"(PARTITION `p20000101` VALUES IN ('20000101'),\n" +
+		"PARTITION `p20220101` VALUES IN ('20220101'),\n" +
+		"PARTITION `p20230516` VALUES IN ('20230516'))")
+	tk.MustExec("analyze table ads_txn")
+	tk.MustExec("set autocommit=OFF;")
+	tk.MustQuery("explain update ads_txn s set aml_cntpr_accno = trim(acpayr_accno) where s._tidb_rowid between 1 and 100000;").Check(testkit.Rows(
+		"Update_5 N/A root  N/A",
+		"└─Projection_6 8000.00 root  test.ads_txn.cusno, test.ads_txn.txn_dt, test.ads_txn.unn_trno, test.ads_txn.aml_cntpr_accno, test.ads_txn.acpayr_accno, test.ads_txn._tidb_rowid",
+		"  └─SelectLock_7 8000.00 root  for update 0",
+		"    └─TableReader_9 10000.00 root partition:all data:TableRangeScan_8",
+		"      └─TableRangeScan_8 10000.00 cop[tikv] table:s range:[1,100000], keep order:false, stats:pseudo"))
+}
+
 func TestIssue45758(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)

planner/core/logical_plans.go

@@ -1447,15 +1447,14 @@ func (ds *DataSource) deriveTablePathStats(path *util.AccessPath, conds []expres
 	path.CountAfterAccess = float64(ds.statisticTable.Count)
 	path.TableFilters = conds
 	var pkCol *expression.Column
-	columnLen := len(ds.schema.Columns)
 	isUnsigned := false
 	if ds.tableInfo.PKIsHandle {
 		if pkColInfo := ds.tableInfo.GetPkColInfo(); pkColInfo != nil {
 			isUnsigned = mysql.HasUnsignedFlag(pkColInfo.GetFlag())
 			pkCol = expression.ColInfo2Col(ds.schema.Columns, pkColInfo)
 		}
-	} else if columnLen > 0 && ds.schema.Columns[columnLen-1].ID == model.ExtraHandleID {
-		pkCol = ds.schema.Columns[columnLen-1]
+	} else {
+		pkCol = ds.schema.GetExtraHandleColumn()
 	}
 	if pkCol == nil {
 		path.Ranges = ranger.FullIntRange(isUnsigned)