WITH ... AS MATERIALIZED returns wrong rows under ORDER BY ... LIMIT: top-k aggregation limits the shared CTE body #27219

Describe the bug

ORDER BY x LIMIT n above a MATERIALIZED CTE whose body groups by x returns wrong rows, often none. The physical optimizer rule TopKAggregation puts the limit on the aggregate inside the CTE body, so the body keeps only n groups, and every reference to the CTE sees only those. With NOT MATERIALIZED the rows are right. The feature (#25580) is on branch-56, so 56.0.0 would ship with this.

To Reproduce

In datafusion-cli (the PostgreSQL dialect is needed to parse MATERIALIZED):

SET datafusion.sql_parser.dialect = 'PostgreSQL';
CREATE TABLE t AS SELECT value AS x FROM generate_series(1, 200);

-- Query 1
SELECT * FROM (
  WITH c AS MATERIALIZED (SELECT x FROM t GROUP BY x)
  SELECT x FROM c WHERE x > 100
) s ORDER BY x LIMIT 3;

-- Query 2
SELECT * FROM (
  WITH c AS MATERIALIZED (SELECT x FROM t GROUP BY x)
  SELECT c.x, m.n FROM c CROSS JOIN (SELECT count(*) AS n FROM c) m
) s ORDER BY x LIMIT 3;

Expected behavior

With NOT MATERIALIZED, query 1 returns 101, 102, 103.

Additional context

EXPLAIN of query 1 shows the limit inside the body:

SortExec: TopK(fetch=3), expr=[x@0 ASC NULLS LAST], preserve_partitioning=[true]
  MaterializedCteExec: name=c
    AggregateExec: mode=FinalPartitioned, gby=[x@0 as x], aggr=[], lim=[3]
      RepartitionExec: partitioning=Hash([x@0], 12), input_partitions=1
        AggregateExec: mode=Partial, gby=[x@0 as x], aggr=[], lim=[3]
          DataSourceExec: partitions=1, partition_sizes=[1]
    FilterExec: x@0 > 100
      MaterializedCteScanExec: name=c

• Cause: TopKAggregation::transform_sort walks down with transform_down (datafusion/physical-optimizer/src/topk_aggregation.rs:91) through every node whose cardinality_effect is Equal (:83). MaterializedCteExec reports Equal (datafusion/physical-plan/src/materialized_cte.rs:623) and lists the body as its first child (:534-536). So the walk enters the body, finds GROUP BY x matching the sort column x, and limits it. The body is not the node's output; the continuation is.
• Fix direction: the walk should only enter children whose rows reach the node's output (for MaterializedCteExec, the continuation).

cc @adriangb (author of #25580)

Willingness to contribute

I'm not planning to fix this myself, so anyone is welcome to pick it up. Feel free to ping me with questions about the cause or when you'd like a review.