TopK aggregation keeps the wrong groups when a projection computes a new column with the aggregate's name #26691

Describe the bug

ORDER BY <col> LIMIT k above an aggregate can be pushed into the aggregate (lim=[k]), so the aggregate keeps only the top k groups. To find the aggregate column, the rule follows the sort column by name through projections. If a projection computes a new value and gives it the same name as the aggregate output, for example -"max(t.v)" AS "max(t.v)", the rule treats it as the aggregate column and picks the top groups by the original value instead of the computed one.

To Reproduce

In datafusion-cli:

CREATE TABLE t AS SELECT column1 AS g, column2 AS v
FROM (VALUES ('a', 2), ('b', 10), ('c', 9), ('d', 3), ('e', 7));

SELECT g, -"max(t.v)" AS "max(t.v)"
FROM (SELECT g, max(v) FROM t GROUP BY g)
ORDER BY 2 DESC LIMIT 2;

+---+----------+
| g | max(t.v) |
+---+----------+
| c | -9       |
| b | -10      |
+---+----------+

SortPreservingMergeExec: [max(t.v)@1 DESC], fetch=2
  SortExec: TopK(fetch=2), expr=[max(t.v)@1 DESC], preserve_partitioning=[true]
    ProjectionExec: expr=[g@0 as g, (- max(t.v)@1) as max(t.v)]
      AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[max(t.v)], lim=[2]
        RepartitionExec: partitioning=Hash([g@0], 12), input_partitions=1
          AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[max(t.v)], lim=[2]

Expected behavior

The two largest values of -max(v) are -2 (a) and -3 (d). That is what SET datafusion.optimizer.enable_topk_aggregation = false; returns:

+---+----------+
| g | max(t.v) |
+---+----------+
| a | -2       |
| d | -3       |
+---+----------+

Additional context

The name tracking is in datafusion/physical-optimizer/src/topk_aggregation.rs (around lines 83-92). For each projection expression, the rule only looks at plain Column expressions (continue on anything else) to follow renames. It never notices that the name it is tracking now belongs to a computed expression, so it keeps looking for that name below and finds the aggregate's max(t.v). A fix is to stop (no TopK aggregation) when the tracked name is produced by any expression other than a plain column.

The same walk continues through any other operator whose cardinality_effect() is Equal or GreaterEqual and keeps the name unchanged (around lines 95-96). That is only correct if the operator also passes the column through unchanged.

The group key has to be a type TopK aggregation supports. With a VARCHAR column (which becomes Utf8View) the rule does not fire, and the result is correct.

Willingness to contribute

I cannot contribute a fix for this bug at this time