TopK queries can fail plan validation after exact sort pushdown #26739

Describe the bug

With sort pushdown enabled, an ORDER BY ... LIMIT query can fail plan validation when a source accepts the requested ordering but cannot apply the row limit itself.

In the standalone sort-pushdown path, removing a partition-preserving TopK (a sort with a row limit) replaces its per-partition limit with GlobalLimitExec. That operator requires a single input partition, so the plan is invalid if the source still has multiple partitions:

Plan: ["GlobalLimitExec: skip=0, fetch=3", "  ExactSortSource"]
does not satisfy distribution requirements: SinglePartition.
Child-0 output partitioning: UnknownPartitioning(4)

To Reproduce

Register a custom source named exact_sort_source that:

• Has four partitions, with values 0..15 interleaved across them. Values 0..7 belong to group 0; values 8..15 belong to group 1.
• Advertises no ordering until try_pushdown_sort is called, then returns SortOrderPushdownResult::Exact for (group_id ASC, value ASC), preserving all four partitions.
• Does not support with_fetch.

The complete source and data setup is included in the regression test. Then run:

SET datafusion.optimizer.enable_sort_pushdown = true;
SET datafusion.execution.target_partitions = 4;

SELECT group_id, value,
       SUM(value) OVER (
           PARTITION BY group_id ORDER BY value ROWS UNBOUNDED PRECEDING
       ) AS running_sum
FROM exact_sort_source
ORDER BY group_id, value
LIMIT 3;

On unfixed code, this fails with the distribution error above.

Expected behavior

The query should execute successfully and return:

group_id  value  running_sum
0         0      0
0         1      1
0         2      3

Removing the sort must retain its per-partition limit, using LocalLimitExec when the source cannot apply it.

Additional context

The window uses native SUM(...) OVER (...). Its PARTITION BY requires redistributing this source's rows by group, and the outer LIMIT lets the optimizer add a local TopK before that redistribution. This reaches the affected standalone path. Without the window, this query takes the already-correct merge/sort path.

A fix and an SLT regression test, including EXPLAIN, are in #26736.

Willingness to contribute

I plan to contribute a fix for this bug independently

Developed with AI assistance.