planner: make queries with the extra column _tidb_rowid can use Point…

…Get (#32705) close #31543
pingcap · Mar 4, 2022 · 95b05d1 · 95b05d1
1 parent f996001
commit 95b05d1
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 0 deletions.
diff --git a/planner/core/find_best_task.go b/planner/core/find_best_task.go
@@ -725,6 +725,10 @@ func (ds *DataSource) getPruningInfo(candidates []*candidatePath, prop *property
 
 func (ds *DataSource) isPointGetConvertableSchema() bool {
 	for _, col := range ds.Columns {
+		if col.Name.L == model.ExtraHandleName.L {
+			continue
+		}
+
 		// Only handle tables that all columns are public.
 		if col.State != model.StatePublic {
 			return false

diff --git a/planner/core/point_get_plan.go b/planner/core/point_get_plan.go
@@ -1561,6 +1561,11 @@ func buildPointDeletePlan(ctx sessionctx.Context, pointPlan PhysicalPlan, dbName
 }
 
 func findCol(tbl *model.TableInfo, colName *ast.ColumnName) *model.ColumnInfo {
+	if colName.Name.L == model.ExtraHandleName.L && !tbl.PKIsHandle {
+		colInfo := model.NewExtraHandleColInfo()
+		colInfo.Offset = len(tbl.Columns) - 1
+		return colInfo
+	}
 	for _, col := range tbl.Columns {
 		if col.Name.L == colName.Name.L {
 			return col

diff --git a/planner/core/point_get_plan_test.go b/planner/core/point_get_plan_test.go
@@ -184,6 +184,37 @@ func TestPointGetForUpdate(t *testing.T) {
 	tk.MustExec("rollback")
 }
 
+func TestGetExtraColumn(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`CREATE TABLE t (
+	  a int(11) DEFAULT NULL,
+	  b int(11) DEFAULT NULL,
+	  UNIQUE KEY idx (a))`)
+	tk.MustQuery(`explain format='brief' select t.*, _tidb_rowid from t where a = 1`).Check(testkit.Rows("Point_Get 1.00 root table:t, index:idx(a) "))
+	tk.MustQuery(`explain format='brief' select t.*, _tidb_rowid, date_format(a, "") from t where a = 1`).Check(testkit.Rows(
+		`Projection 1.00 root  test.t.a, test.t.b, test.t._tidb_rowid, date_format(cast(test.t.a, datetime BINARY), )->Column#4`,
+		`└─Point_Get 1.00 root table:t, index:idx(a) `))
+	tk.MustExec(`begin`) // in transaction
+	tk.MustExec(`insert into t values (1, 1)`)
+	tk.MustQuery(`explain format='brief' select t.*, _tidb_rowid from t where a = 1`).Check(testkit.Rows(`Point_Get 1.00 root table:t, index:idx(a) `))
+	tk.MustExec(`commit`)
+	tk.MustQuery(`explain format='brief' select count(_tidb_rowid) from t where a=1`).Check(testkit.Rows(
+		`StreamAgg 1.00 root  funcs:count(test.t._tidb_rowid)->Column#4`,
+		`└─Point_Get 1.00 root table:t, index:idx(a) `))
+	tk.MustQuery(`explain format='brief' select *, date_format(b, "") from t where a =1 for update`).Check(testkit.Rows(
+		`Projection 1.00 root  test.t.a, test.t.b, date_format(cast(test.t.b, datetime BINARY), )->Column#4`,
+		`└─SelectLock 1.00 root  for update 0`,
+		`  └─Point_Get 1.00 root table:t, index:idx(a) `))
+
+	// if the PK is handled
+	tk.MustExec(`create table t1 (pk int, a int, b int, primary key(pk), unique key(a))`)
+	err := tk.ExecToErr(`explain format='brief' select t1.*, _tidb_rowid from t1 where a = 1`)
+	require.EqualError(t, err, `[planner:1054]Unknown column '_tidb_rowid' in 'field list'`)
+}
+
 func TestPointGetForUpdateWithSubquery(t *testing.T) {
 	store, clean := testkit.CreateMockStore(t)
 	defer clean()