fix: skip TopK aggregation when sort column is computed

datafusion/physical-optimizer/src/topk_aggregation.rs

@@ -68,14 +68,23 @@ impl TopKAggregation {
                     Some(aggr) => return Ok(Transformed::yes(Arc::new(aggr.data))),
                 }
             } else if let Some(proj) = plan.downcast_ref::<ProjectionExec>() {
-                // track renames due to successive projections
+                // Follow the sort column through renames. If that output is a
+                // computed expression (for example `-max(v) AS "max(v)"`), the
+                // aggregate below no longer produces the values we sort by, so
+                // skip TopK aggregation for this sort.
                 for proj_expr in proj.expr() {
-                    let Some(src_col) = proj_expr.expr.downcast_ref::<Column>() else {
+                    if proj_expr.alias != sort_col_name {
                         continue;
-                    };
-                    if proj_expr.alias == sort_col_name {
-                        sort_col_name = src_col.name().to_string();
                     }
+                    match proj_expr.expr.downcast_ref::<Column>() {
+                        Some(src_col) => {
+                            sort_col_name = src_col.name().to_string();
+                        }
+                        None => {
+                            cardinality_preserved = false;
+                        }
+                    }
+                    break;
                 }
             } else {
                 // or we continue down through types that don't reduce cardinality

datafusion/sqllogictest/test_files/aggregates_topk.slt

@@ -886,3 +886,65 @@ GROUP BY k ORDER BY MAX(v) DESC NULLS LAST LIMIT 1;
 
 statement ok
 set datafusion.execution.target_partitions = 4;
+
+# Regression for #26691: a projection that computes a new value under the
+# aggregate's output name must not keep TopK aggregation on the aggregate.
+# Ordering by -max(v) DESC must keep the groups with the smallest maxima.
+statement ok
+CREATE TABLE topk_negated_max AS SELECT column1 AS g, column2 AS v
+FROM (VALUES ('a', 2), ('b', 10), ('c', 9), ('d', 3), ('e', 7));
+
+statement ok
+set datafusion.optimizer.enable_topk_aggregation = true;
+
+query TI
+SELECT g, -"max(topk_negated_max.v)" AS "max(topk_negated_max.v)"
+FROM (SELECT g, max(v) FROM topk_negated_max GROUP BY g)
+ORDER BY 2 DESC LIMIT 2;
+----
+a -2
+d -3
+
+# TopK aggregation must not apply: lim would keep the largest maxima first.
+query TT
+EXPLAIN SELECT g, -"max(topk_negated_max.v)" AS "max(topk_negated_max.v)"
+FROM (SELECT g, max(v) FROM topk_negated_max GROUP BY g)
+ORDER BY 2 DESC LIMIT 2;
+----
+logical_plan
+01)Sort: max(topk_negated_max.v) DESC NULLS FIRST, fetch=2
+02)--Projection: topk_negated_max.g, (- max(topk_negated_max.v)) AS max(topk_negated_max.v)
+03)----Aggregate: groupBy=[[topk_negated_max.g]], aggr=[[max(topk_negated_max.v)]]
+04)------TableScan: topk_negated_max projection=[g, v]
+physical_plan
+01)SortPreservingMergeExec: [max(topk_negated_max.v)@1 DESC], fetch=2
+02)--SortExec: TopK(fetch=2), expr=[max(topk_negated_max.v)@1 DESC], preserve_partitioning=[true]
+03)----ProjectionExec: expr=[g@0 as g, (- max(topk_negated_max.v)@1) as max(topk_negated_max.v)]
+04)------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[max(topk_negated_max.v)]
+05)--------RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1
+06)----------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[max(topk_negated_max.v)]
+07)------------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# Sibling computed column: sorting by the plain aggregate still allows TopK.
+query TT
+EXPLAIN SELECT g, max(v), -max(v) AS neg
+FROM topk_negated_max
+GROUP BY g
+ORDER BY max(v) DESC LIMIT 2;
+----
+logical_plan
+01)Sort: max(topk_negated_max.v) DESC NULLS FIRST, fetch=2
+02)--Projection: topk_negated_max.g, max(topk_negated_max.v), (- max(topk_negated_max.v)) AS neg
+03)----Aggregate: groupBy=[[topk_negated_max.g]], aggr=[[max(topk_negated_max.v)]]
+04)------TableScan: topk_negated_max projection=[g, v]
+physical_plan
+01)SortPreservingMergeExec: [max(topk_negated_max.v)@1 DESC], fetch=2
+02)--SortExec: TopK(fetch=2), expr=[max(topk_negated_max.v)@1 DESC], preserve_partitioning=[true]
+03)----ProjectionExec: expr=[g@0 as g, max(topk_negated_max.v)@1 as max(topk_negated_max.v), (- max(topk_negated_max.v)@1) as neg]
+04)------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[max(topk_negated_max.v)], lim=[2]
+05)--------RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1
+06)----------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[max(topk_negated_max.v)], lim=[2]
+07)------------DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+DROP TABLE topk_negated_max;