select failed for partition table with generated column with error `needPartitionHandle != ret, tp(1)`

[joechenrh]

Bug Report

Please answer these questions before submitting your issue. Thanks!

1. Minimal reproduce step (Required)

drop table tp;

CREATE TABLE tp (
  id int, c1 json, c2 json GENERATED ALWAYS AS (c1) VIRTUAL, KEY `idx_31` (id)
) PARTITION BY RANGE (id)
(PARTITION `p0` VALUES LESS THAN (2),
 PARTITION `p1` VALUES LESS THAN (4532022));

INSERT INTO tp (id, c1) VALUES (0, "[1]");

select id from tp where json_contains(c2, "1") group by id having id in (0);

2. What did you expect to see? (Required)

+------+                                                                                                                                                                                                         
| id   |                                                                                                                                                                                                         
+------+                                                                                                                                                                                                         
|    0 |                                                                                                                                                                                                         
+------+ 

3. What did you see instead (Required)

ERROR 1105 (HY000): Internal error, needPartitionHandle != ret, tp(1)

4. What is your TiDB version? (Required)

[joechenrh]

diff --git a/pkg/planner/core/physical_plans.go b/pkg/planner/core/physical_plans.go
index 1887ea8d39..cf5fa3b9a5 100644
--- a/pkg/planner/core/physical_plans.go
+++ b/pkg/planner/core/physical_plans.go
@@ -1061,14 +1061,25 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema,
 	colsInfo []*model.ColumnInfo) []*model.ColumnInfo {
 	copyColumn := make([]*model.ColumnInfo, len(columns))
 	copy(copyColumn, columns)
-	var extraColumn *expression.Column
-	var extraColumnModel *model.ColumnInfo
-	if schema.Columns[len(schema.Columns)-1].ID == model.ExtraHandleID {
-		extraColumn = schema.Columns[len(schema.Columns)-1]
-		extraColumnModel = copyColumn[len(copyColumn)-1]
-		schema.Columns = schema.Columns[:len(schema.Columns)-1]
-		copyColumn = copyColumn[:len(copyColumn)-1]
+
+	oldNumColumns := len(schema.Columns)
+	numExtraColumns := 0
+	for i := oldNumColumns - 1; i >= 0; i-- {
+		cid := schema.Columns[i].ID
+		if cid != model.ExtraHandleID && cid != model.ExtraPhysTblID {
+			break
+		}
+		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]
+
 	schemaColumns := schema.Columns
 	for _, col := range schemaColumns {
 		if col.VirtualExpr == nil {
@@ -1083,10 +1094,9 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema,
 			}
 		}
 	}
-	if extraColumn != nil {
-		schema.Columns = append(schema.Columns, extraColumn)
-		copyColumn = append(copyColumn, extraColumnModel) // nozero
-	}
+
+	schema.Columns = append(schema.Columns, extraColumns...)
+	copyColumn = append(copyColumn, extraColumnModels...)
 	return copyColumn
 }

Apply this patch can solve the problems in the issue, but I'm not sure whether its the correct solution.

[Defined2014]

The diff patch looks great, I think you could push a PR to fix this issue.