planner: fix column pruner will clear handleCols and lead filling _tidb_rowid when datasource's schema length is 0 for a pkIsHandled table (#45217)

planner/core/handle_cols.go

@@ -46,7 +46,7 @@ type HandleCols interface {
 	BuildPartitionHandleFromIndexRow(row chunk.Row) (kv.PartitionHandle, error)
 	// ResolveIndices resolves handle column indices.
 	ResolveIndices(schema *expression.Schema) (HandleCols, error)
-	// IsInt returns if the HandleCols is a single tnt column.
+	// IsInt returns if the HandleCols is a single int column.
 	IsInt() bool
 	// String implements the fmt.Stringer interface.
 	String() string

planner/core/logical_plan_builder.go

@@ -4789,6 +4789,7 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as
 		}
 	}
 	ds.handleCols = handleCols
+	ds.unMutableHandleCols = handleCols
 	handleMap := make(map[int64][]HandleCols)
 	handleMap[tableInfo.ID] = []HandleCols{handleCols}
 	b.handleHelper.pushMap(handleMap)

planner/core/logical_plans.go

@@ -1210,7 +1210,8 @@ type DataSource struct {
 	// handleCol represents the handle column for the datasource, either the
 	// int primary key column or extra handle column.
 	// handleCol *expression.Column
-	handleCols HandleCols
+	handleCols          HandleCols
+	unMutableHandleCols HandleCols
 	// TblCols contains the original columns of table before being pruned, and it
 	// is used for estimating table scan cost.
 	TblCols []*expression.Column

planner/core/rule_column_pruning.go

@@ -343,6 +343,11 @@ func (ds *DataSource) PruneColumns(parentUsedCols []*expression.Column, opt *log
 		ds.Columns = append(ds.Columns, handleColInfo)
 		ds.schema.Append(handleCol)
 	}
+	// ref: https://github.com/pingcap/tidb/issues/44579
+	// when first entering columnPruner, we kept a column-a in datasource since upper agg function count(a) is used.
+	//		then we mark the handleCols as nil here.
+	// when second entering columnPruner, the count(a) is eliminated since it always not null. we should fill another
+	// 		extra col, in this way, handle col is useful again, otherwise, _tidb_rowid will be filled.
 	if ds.handleCols != nil && ds.handleCols.IsInt() && ds.schema.ColumnIndex(ds.handleCols.GetCol(0)) == -1 {
 		ds.handleCols = nil
 	}
@@ -659,6 +664,11 @@ func preferKeyColumnFromTable(dataSource *DataSource, originColumns []*expressio
 		if dataSource.handleCols != nil {
 			resultColumn = dataSource.handleCols.GetCol(0)
 			resultColumnInfo = resultColumn.ToInfo()
+		} else if dataSource.table.Meta().PKIsHandle {
+			// dataSource.handleCols = nil doesn't mean datasource doesn't have a intPk handle.
+			// since datasource.handleCols will be cleared in the first columnPruner.
+			resultColumn = dataSource.unMutableHandleCols.GetCol(0)
+			resultColumnInfo = resultColumn.ToInfo()
 		} else {
 			resultColumn = dataSource.newExtraHandleSchemaCol()
 			resultColumnInfo = model.NewExtraHandleColInfo()