Repository: pingcap/tidb
File: pkg/planner/core/operator/physicalop/physical_table_scan.go (Go)

Code before:
		pkCol := expression.ColInfo2Col(ts.TblCols, pk)
		if !ts.Schema().Contains(pkCol) {
			ts.Schema().Append(pkCol)
			ts.Columns = append(ts.Columns, pk)
			columnAdded = true
		}
	} else if ts.Table.IsCommonHandle {
		idxInfo := ts.Table.GetPrimaryKey()
		for _, idxCol := range idxInfo.Columns {
			col := ts.TblCols[idxCol.Offset]
			if !ts.Schema().Contains(col) {
				columnAdded = true
				ts.Schema().Append(col)
				ts.Columns = append(ts.Columns, col.ToInfo())
			}
		}
	} else if ts.HandleCols != nil && !ts.Schema().Contains(ts.HandleCols.GetCol(0)) {
		// For cluster tables, HandleCols is nil and we skip adding ExtraHandleID.
		ts.Schema().Append(ts.HandleCols.GetCol(0))
		ts.Columns = append(ts.Columns, model.NewExtraHandleColInfo())
		columnAdded = true
	}

	// For the global index of the partitioned table, we also need the PhysicalTblID to identify the rows from each partition.
	if ts.Table.GetPartitionInfo() != nil && ts.SCtx().GetSessionVars().StmtCtx.UseDynamicPartitionPrune() {
		tmpColumns, tmpSchema, newColAdded := AddExtraPhysTblIDColumn(ts.SCtx(), ts.Columns, ts.Schema())
		ts.Columns = tmpColumns
		ts.SetSchema(tmpSchema)
		columnAdded = columnAdded || newColAdded
	}
	return currentTopPlan, nil, columnAdded, nil
}

// extractFiltersForIndexMerge returns:
// `pushed`: exprs that can be pushed to TiKV.
// `remaining`: exprs that can NOT be pushed to TiKV but can be pushed to other storage engines.
// Why do we need this func?
// IndexMerge only works on TiKV, so we need to find all exprs that cannot be pushed to TiKV, and add a new Selection above IndexMergeReader.
//
//	But the new Selection should exclude the exprs that can NOT be pushed to ALL the storage engines.
//	Because these exprs have already been put in another Selection(check rule_predicate_push_down).
func extractFiltersForIndexMerge(ctx expression.PushDownContext, filters []expression.Expression) (pushed []expression.Expression, remaining []expression.Expression) {
	for _, expr := range filters {
		if expression.CanExprsPushDown(ctx, []expression.Expression{expr}, kv.TiKV) {
			pushed = append(pushed, expr)
			continue
		}
		if expression.CanExprsPushDown(ctx, []expression.Expression{expr}, kv.UnSpecified) {
			remaining = append(remaining, expr)
		}
	}

Review comment (performance):
currently we rely on preferKeyColumnFromTable to supple the additional handle for cluster table without any handle cols, i think we can just also add that logic in here too.

Otherwise, when optimizer block the column-prune-logical rule, the clusters can never have has its handle cols themselves when its nil here.

Code after:
	// If we don't need to use ordered scan, we don't need do the following codes for adding new columns.
	if !matchProp {
		return currentTopPlan, nil, false, nil
	}

	// Add the row handle into the schema.
	columnAdded := false
	if ts.Table.PKIsHandle {
		pk := ts.Table.GetPkColInfo()
		pkCol := expression.ColInfo2Col(ts.TblCols, pk)
		if !ts.Schema().Contains(pkCol) {
			ts.Schema().Append(pkCol)
			ts.Columns = append(ts.Columns, pk)
			columnAdded = true
		}
	} else if ts.Table.IsCommonHandle {
		idxInfo := ts.Table.GetPrimaryKey()
		for _, idxCol := range idxInfo.Columns {
			col := ts.TblCols[idxCol.Offset]
			if !ts.Schema().Contains(col) {
				columnAdded = true
				ts.Schema().Append(col)
				ts.Columns = append(ts.Columns, col.ToInfo())
			}
		}
	} else if ts.HandleCols != nil && !ts.Schema().Contains(ts.HandleCols.GetCol(0)) {
		ts.Schema().Append(ts.HandleCols.GetCol(0))
		ts.Columns = append(ts.Columns, model.NewExtraHandleColInfo())
		columnAdded = true
	} else if ds.Table.Type().IsClusterTable() {
		// For cluster tables without HandleCols, use the first column like preferKeyColumnFromTable does.
		// This handles the case when column-prune-logical rule is blocked and cluster tables don't have
		// their handle cols set up through the normal column pruning path.
		if len(ts.TblCols) > 0 {
			col := ts.TblCols[0]
			if !ts.Schema().Contains(col) {
				ts.Schema().Append(col)
				// Find the corresponding ColumnInfo from table metadata
				if colInfo := model.FindColumnInfoByID(ts.Table.Columns, col.ID); colInfo != nil {
					ts.Columns = append(ts.Columns, colInfo)
				} else {
					ts.Columns = append(ts.Columns, col.ToInfo())
				}
				columnAdded = true
			}
		} else if len(ts.Table.Columns) > 0 {
			// Fallback to table metadata columns
			colMeta := ts.Table.Columns[0]
			col := &expression.Column{
				RetType:  colMeta.FieldType.Clone(),
				UniqueID: ts.SCtx().GetSessionVars().AllocPlanColumnID(),