planner: Handle _tidb_rowid correctly in batchPointGet Plan to avoid index out of range error (#58687) (#60441)

close #58581

Author: ti-chi-bot

pkg/planner/core/find_best_task.go

@@ -1493,7 +1493,7 @@ func findBestTask4LogicalDataSource(lp base.LogicalPlan, prop *property.Physical
 			}
 			if canConvertPointGet && len(path.Ranges) > 1 {
 				// TODO: This is now implemented, but to decrease
-				// the impact of supporting plan cache for patitioning,
+				// the impact of supporting plan cache for partitioning,
 				// this is not yet enabled.
 				// TODO: just remove this if block and update/add tests...
 				// We can only build batch point get for hash partitions on a simple column now. This is

pkg/planner/core/physical_plans.go

@@ -1059,25 +1059,36 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema,
 
 	oldNumColumns := len(schema.Columns)
 	numExtraColumns := 0
+	ordinaryColumnExists := false
 	for i := oldNumColumns - 1; i >= 0; i-- {
 		cid := schema.Columns[i].ID
 		// Move extra columns to the end.
 		// ExtraRowChecksumID is ignored here since it's treated as an ordinary column.
 		// https://github.com/pingcap/tidb/blob/3c407312a986327bc4876920e70fdd6841b8365f/pkg/util/rowcodec/decoder.go#L206-L222
 		if cid != model.ExtraHandleID && cid != model.ExtraPhysTblID {
+			ordinaryColumnExists = true
 			break
 		}
 		numExtraColumns++
 	}
+	if ordinaryColumnExists && numExtraColumns > 0 {
+		extraColumns := make([]*expression.Column, numExtraColumns)
+		copy(extraColumns, schema.Columns[oldNumColumns-numExtraColumns:])
+		schema.Columns = schema.Columns[:oldNumColumns-numExtraColumns]
 
-	extraColumns := make([]*expression.Column, numExtraColumns)
-	copy(extraColumns, schema.Columns[oldNumColumns-numExtraColumns:])
-	schema.Columns = schema.Columns[:oldNumColumns-numExtraColumns]
+		extraColumnModels := make([]*model.ColumnInfo, numExtraColumns)
+		copy(extraColumnModels, copyColumn[len(copyColumn)-numExtraColumns:])
+		copyColumn = copyColumn[:len(copyColumn)-numExtraColumns]
 
-	extraColumnModels := make([]*model.ColumnInfo, numExtraColumns)
-	copy(extraColumnModels, copyColumn[len(copyColumn)-numExtraColumns:])
-	copyColumn = copyColumn[:len(copyColumn)-numExtraColumns]
+		copyColumn = expandVirtualColumn(schema, copyColumn, colsInfo)
+		schema.Columns = append(schema.Columns, extraColumns...)
+		copyColumn = append(copyColumn, extraColumnModels...)
+		return copyColumn
+	}
+	return expandVirtualColumn(schema, copyColumn, colsInfo)
+}
 
+func expandVirtualColumn(schema *expression.Schema, copyColumn []*model.ColumnInfo, colsInfo []*model.ColumnInfo) []*model.ColumnInfo {
 	schemaColumns := schema.Columns
 	for _, col := range schemaColumns {
 		if col.VirtualExpr == nil {
@@ -1092,9 +1103,6 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema,
 			}
 		}
 	}
-
-	schema.Columns = append(schema.Columns, extraColumns...)
-	copyColumn = append(copyColumn, extraColumnModels...)
 	return copyColumn
 }
 

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

@@ -815,3 +815,13 @@ Projection_9	12.50	root		test.t.a, test.t1.a, test.t1.b
 select /*+ inl_join(t1), use_index(t1, idx) */ * from t join t1 on t.a = t1.a and t1.b = 123;
 a	a	b
 2	2	123
+drop table if exists t;
+create table t (id int unique key, c int);
+insert into t values (1, 10);
+insert into t values (2, 20);
+insert into t values (3, 30);
+select _tidb_rowid from t where id in (1, 2, 3);
+_tidb_rowid
+1
+2
+3

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

@@ -553,3 +553,11 @@ insert into t1 values(2, 123), (123, 2);
 set tidb_opt_fix_control='44855:on';
 explain select /*+ inl_join(t1), use_index(t1, idx) */ * from t join t1 on t.a = t1.a and t1.b = 123;
 select /*+ inl_join(t1), use_index(t1, idx) */ * from t join t1 on t.a = t1.a and t1.b = 123;
+
+# TestIssue58581
+drop table if exists t;
+create table t (id int unique key, c int);
+insert into t values (1, 10);
+insert into t values (2, 20);
+insert into t values (3, 30);
+select _tidb_rowid from t where id in (1, 2, 3);