planner: Handle _tidb_rowid correctly in batchPointGet Plan to avoid …

…index out of range error (#58687) (#59175)

close #58581

planner/core/integration_partition_test.go

@@ -1693,3 +1693,17 @@ func TestGeneratedColumnWithPartition(t *testing.T) {
 	tk.MustExec(`INSERT INTO tp (id, c1) VALUES (0, 1)`)
 	tk.MustExec(`select /*+ FORCE_INDEX(tp, idx) */id from tp where c2 = 2 group by id having id in (0)`)
 }
+
+// TestIssue58581
+func TestIssue58581(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t (id int unique key, c int);")
+	tk.MustExec("insert into t values (1, 10);")
+	tk.MustExec("insert into t values (2, 20);")
+	tk.MustExec("insert into t values (3, 30);")
+	tk.MustQuery("select _tidb_rowid from t where id in (1, 2, 3);").Check(testkit.Rows("1", "2", "3"))
+}

planner/core/physical_plans.go

@@ -923,25 +923,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 {
@@ -956,9 +967,6 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema,
 			}
 		}
 	}
-
-	schema.Columns = append(schema.Columns, extraColumns...)
-	copyColumn = append(copyColumn, extraColumnModels...)
 	return copyColumn
 }