Incorrect results: ORDER BY b, a over a subquery with ORDER BY b LIMIT n is not sorted by (b, a) when SortPreservingMergeExec has a fetch #26162

Describe the bug

ORDER BY b, a on top of a subquery with ORDER BY b LIMIT n returns rows that are not sorted by (b, a) when the input has multiple partitions.

The sort pushdown optimizer pushes the finer [b, a] requirement below a SortPreservingMergeExec that has a fetch, but the merge itself keeps its original key [b]. The resulting plan has no operator that actually orders the output by (b, a).

Found by @jayzhan211 while reviewing #25408 (see #25408 (comment)). The bug is not introduced by that PR; it reproduces on main (754d3c1) as well. That comment also discusses the root cause and possible fixes.

To Reproduce

Run the following in datafusion-cli:

set datafusion.execution.target_partitions = 4;
CREATE TABLE src (a INT, b INT) AS VALUES (1, 3), (2, 1), (3, 2), (4, 3), (5, 1), (6, 2), (7, 3), (8, 1);
CREATE TABLE t AS SELECT a, b FROM src UNION ALL SELECT a + 10, b FROM src
  UNION ALL SELECT a + 20, b FROM src UNION ALL SELECT a + 30, b FROM src;
SELECT * FROM (SELECT a, b FROM t ORDER BY b LIMIT 10) ORDER BY b, a;

Actual output (rows are not ordered by a within b = 1):

+----+---+
| a  | b |
+----+---+
| 2  | 1 |
| 22 | 1 |
| 5  | 1 |
| 25 | 1 |
| 8  | 1 |
| 12 | 1 |
| 15 | 1 |
| 18 | 1 |
| 28 | 1 |
| 32 | 1 |
+----+---+

Plan:

EXPLAIN SELECT * FROM (SELECT a, b FROM t ORDER BY b LIMIT 10) ORDER BY b, a;

logical_plan  | Sort: t.b ASC NULLS LAST, t.a ASC NULLS LAST
              |   Sort: t.b ASC NULLS LAST, fetch=10
              |     TableScan: t projection=[a, b]
physical_plan | SortPreservingMergeExec: [b@1 ASC NULLS LAST], fetch=10
              |   SortExec: TopK(fetch=10), expr=[b@1 ASC NULLS LAST, a@0 ASC NULLS LAST], preserve_partitioning=[true]
              |     DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1]

Note the SortPreservingMergeExec merges on [b] only, while the SortExec below it was rewritten to [b, a].

Expected behavior

The output should be ordered by (b, a):

+----+---+
| a  | b |
+----+---+
| 2  | 1 |
| 5  | 1 |
| 8  | 1 |
| 12 | 1 |
| 15 | 1 |
| 18 | 1 |
| 22 | 1 |
| 25 | 1 |
| 28 | 1 |
| 32 | 1 |
+----+---+

Additional context

• Root cause analysis and suggested fixes: fix: Avoid pushing down sort under limits #25408 (comment)
• Related PR: fix: Avoid pushing down sort under limits #25408
• Relevant code: datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs