fix: preserve sort column identity in TopK aggregation

datafusion/physical-optimizer/src/topk_aggregation.rs

@@ -27,6 +27,7 @@ use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::aggregates::AggregateExec;
 use datafusion_physical_plan::execution_plan::CardinalityEffect;
+use datafusion_physical_plan::materialized_cte::MaterializedCteExec;
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
 use itertools::Itertools;
@@ -64,7 +65,7 @@ impl TopKAggregation {
         let order = order.iter().exactly_one().ok()?;
         let sort_options = order.options;
         let order = order.expr.downcast_ref::<Column>()?;
-        let mut sort_col_name = order.name().to_string();
+        let mut sort_col = order.clone();
         let limit = sort.fetch()?;
 
         let mut cardinality_preserved = true;
@@ -75,20 +76,24 @@ impl TopKAggregation {
             if let Some(aggr) = plan.downcast_ref::<AggregateExec>() {
                 match aggr
                     .clone()
-                    .try_optimize_topk(limit, &sort_col_name, sort_options)
+                    .try_optimize_topk(limit, sort_col.name(), sort_options)
                 {
                     None => cardinality_preserved = false,
                     Some(aggr) => return Ok(Transformed::yes(Arc::new(aggr.data))),
                 }
+            } else if plan.downcast_ref::<MaterializedCteExec>().is_some() {
+                // Its body produces rows; the output schema comes from the continuation.
+                cardinality_preserved = false;
             } else if let Some(proj) = plan.downcast_ref::<ProjectionExec>() {
-                // track renames due to successive projections
-                for proj_expr in proj.expr() {
-                    let Some(src_col) = proj_expr.expr.downcast_ref::<Column>() else {
-                        continue;
-                    };
-                    if proj_expr.alias == sort_col_name {
-                        sort_col_name = src_col.name().to_string();
-                    }
+                // Only follow direct columns: computed expressions can change group ranking.
+                let src_col = proj
+                    .expr()
+                    .get(sort_col.index())
+                    .and_then(|proj_expr| proj_expr.expr.downcast_ref::<Column>());
+                if let Some(src_col) = src_col {
+                    sort_col = src_col.clone();
+                } else {
+                    cardinality_preserved = false;
                 }
             } else {
                 // or we continue down through types that don't reduce cardinality

datafusion/sqllogictest/test_files/aggregates_topk.slt

@@ -195,6 +195,80 @@ a -1 -1
 NULL 0 0
 a 1 1
 
+# Computed projections can change which groups belong in the top K, even when
+# their output names match the original aggregate columns (issue #26691).
+statement ok
+CREATE TABLE topk_projection AS
+SELECT column1 AS g, column2 AS v
+FROM (VALUES ('a', 2), ('b', 10), ('c', 9), ('d', 3), ('e', 7));
+
+query TI
+SELECT g, -"max(topk_projection.v)" AS "max(topk_projection.v)"
+FROM (SELECT g, max(v) FROM topk_projection GROUP BY g)
+ORDER BY 2 DESC LIMIT 2;
+----
+a -2
+d -3
+
+# The same ordering change also applies to TopK distinct aggregation.
+query I
+SELECT -v AS v FROM (SELECT DISTINCT v FROM topk_projection)
+ORDER BY v LIMIT 2;
+----
+-10
+-9
+
+# A materialized CTE produces rows from its continuation, whose column positions
+# can differ from those of the CTE body. Ordering by g must not rank by max(v).
+statement ok
+SET datafusion.sql_parser.dialect = 'PostgreSQL';
+
+statement ok
+CREATE TABLE topk_cte AS
+SELECT column1 AS g, column2 AS v
+FROM (VALUES (1, 100), (2, 20), (3, 30), (4, 40), (5, 50));
+
+query II
+SELECT q.g, q.score FROM (
+  WITH c AS MATERIALIZED (SELECT max(v) AS score, g FROM topk_cte GROUP BY g)
+  SELECT g, score FROM c
+) q ORDER BY q.g DESC NULLS LAST LIMIT 2;
+----
+5 50
+4 40
+
+statement ok
+SET datafusion.sql_parser.dialect = 'Generic';
+
+# A direct sort column remains eligible through renaming and reordering, even
+# when another projected column is computed.
+query ITI
+SELECT score AS rank, g AS label, -score AS decoration
+FROM (SELECT g, max(v) AS score FROM topk_projection GROUP BY g)
+ORDER BY 1 DESC LIMIT 2;
+----
+10 b -10
+9 c -9
+
+query TT
+EXPLAIN SELECT score AS rank, g AS label, -score AS decoration
+FROM (SELECT g, max(v) AS score FROM topk_projection GROUP BY g)
+ORDER BY 1 DESC LIMIT 2;
+----
+logical_plan
+01)Sort: rank DESC NULLS FIRST, fetch=2
+02)--Projection: max(topk_projection.v) AS rank, topk_projection.g AS label, (- max(topk_projection.v)) AS decoration
+03)----Aggregate: groupBy=[[topk_projection.g]], aggr=[[max(topk_projection.v)]]
+04)------TableScan: topk_projection projection=[g, v]
+physical_plan
+01)SortPreservingMergeExec: [rank@0 DESC], fetch=2
+02)--SortExec: TopK(fetch=2), expr=[rank@0 DESC], preserve_partitioning=[true]
+03)----ProjectionExec: expr=[max(topk_projection.v)@1 as rank, g@0 as label, (- max(topk_projection.v)@1) as decoration]
+04)------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[max(topk_projection.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_projection.v)], lim=[2]
+07)------------DataSourceExec: partitions=1, partition_sizes=[1]
+
 statement ok
 CREATE TABLE string_topk(category varchar, val varchar) AS VALUES
 ('x', 'apple'),