Skip to content

fix: decline sort pushdown through SortPreservingMergeExec with fetch (#26162) - #27060

Merged
jayzhan211 merged 3 commits into
apache:mainfrom
galaxytemple:fix-spm-fetch-sort-pushdown
Nov 8, 2026
Merged

jayzhan211 merged 3 commits into
apache:mainfrom
galaxytemple:fix-spm-fetch-sort-pushdown

Conversation

@galaxytemple

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

In #26162, running ORDER BY b, a over a multi-partition subquery with ORDER BY b LIMIT n produces incorrect ordering on a.

Pushing [b, a] down removed the outer SortExec, but SPM only merges on [b]. As a result, a was not sorted in the final output.

In the discussion, @jayzhan211 suggested two possible approaches: declining pushdown for SPM with fetch, or rewriting SPM's expressions when pushing through it.

I chose Ok(None) (declining pushdown) because:

  • The subquery did not explicitly request the parent's sort. When a limit is present, forcing that sort on the entire dataset across partitions before a limit is unnecessary work.
  • Since pushdown_requirement_to_children only adds requirements to children, I think a logic to rewrite SPM adds extra complexity.

What changes are included in this PR?

  • Return Ok(None) from pushdown_requirement_to_children when plan is SortPreservingMergeExec and has a fetch.
  • Add an SLT regression query in limit.slt.

What is the testing strategy for this PR?

Are there any user-facing changes?

Yes, bug fix: queries like ORDER BY b, a over a multi-partition subquery with ORDER BY b LIMIT n now return correctly ordered results on a. No API changes.

…apache#26162)

SortPreservingMergeExec with fetch merges partitions using only its own
keys, so pushing finer parent requirements down drops the needed outer
sort. Declining pushdown keeps the parent SortExec to correctly order
the final fetched rows.
@github-actions github-actions Bot added optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) labels Nov 3, 2026
@galaxytemple galaxytemple changed the title fix: decline sort pushdown through SortPreservingMergeExec with fetch… fix: decline sort pushdown through SortPreservingMergeExec with fetch (#26162) Nov 3, 2026
@galaxytemple
galaxytemple marked this pull request as ready for review November 3, 2026 21:55
@galaxytemple

Copy link
Copy Markdown
Contributor Author

While working on this, I also noticed that SPM without limit (at L589 else if is_sort_preserving_merge(plan)) has a similar issue. I couldn't find a query that actually reaches that branch, since without Limit and offset EliminateSort drops the child sort, and with offset, GlobalLimitExec block the pushdown earlier.

But in case a query reaches there, pushing down without updating SPM's own merge keys would likely produce incorrect results.(similar to #26654, SPM only preserves its own merge keys across partitions). But since this is related to the performance issue #16217, I thought it's better to open a separate issue to discuss it.

@jayzhan211
jayzhan211 self-requested a review November 7, 2026 02:02

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@galaxytemple thanks for the fix. It returns correct results on every query shape I tried, including a window, an ordered aggregate and a UNION above the merge. I'd like the merge without a fetch handled here too, using the same rule. The three test comments are optional.

@@ -525,20 +525,13 @@ fn pushdown_requirement_to_children(
Ok(None)
}
} else if plan.is::<SortPreservingMergeExec>() && plan.fetch().is_some() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please apply this rule to the merge without a fetch in this PR as well, the case you raised in your comment. The is_sort_preserving_merge arm further down still pushes [a, b] through a merge on [a] and drops the outer sort. I couldn't reach it from SQL either, only by calling pushdown_sorts directly, but one arm can cover both cases and remove 15 lines. With the change below, the optimizer tests and all sqllogictest files pass with no expected plan changing.

 use crate::utils::{
-    add_sort_above_with_distribution, is_sort, is_sort_preserving_merge, is_union,
-    is_window,
+    add_sort_above_with_distribution, is_sort, is_union, is_window,
 };
@@
-    } else if plan.is::<SortPreservingMergeExec>() && plan.fetch().is_some() {
-        // Why we decline pushdown here:
-        // 1. SPM merges using only its own keys, so pushing additional
-        //    parent keys down cannot be preserved anyway.
-        // 2. The subquery did not explicitly request the parent's sort, so forcing
-        //    it on the entire dataset before a fetch is an unnecessary cost.
-        // 3. Decline pushdown so the parent sort is not dropped and orders the fetched rows.
+    } else if plan.is::<SortPreservingMergeExec>() {
+        // A merge orders rows across partitions only by its own keys, and this
+        // function doesn't rewrite them. Keep the finer parent sort above it.
         Ok(None)
@@
-    } else if is_sort_preserving_merge(plan) {
-        let new_ordering = LexOrdering::from(parent_required.first().clone());
-        let mut spm_eqs = plan.equivalence_properties().clone();
-        let old_ordering = spm_eqs.output_ordering().unwrap();
-        // Sort preserving merge will have new ordering, one requirement above is pushed down to its below.
-        let change = spm_eqs.reorder(new_ordering)?;
-        if !change || spm_eqs.ordering_satisfy(old_ordering)? {
-            // Can push-down through SortPreservingMergeExec, because parent requirement is finer
-            // than SortPreservingMergeExec output ordering.
-            Ok(Some(vec![Some(parent_required)]))
-        } else {
-            // Do not push-down through SortPreservingMergeExec when
-            // ordering requirement invalidates requirement of sort preserving merge exec.
-            Ok(None)
-        }
     } else if let Some(hash_join) = plan.downcast_ref::<HashJoinExec>() {
Test that fails without the change

For datafusion/core/tests/physical_optimizer/ensure_requirements.rs:

/// A merge guarantees only its own keys across partitions, so a finer
/// requirement must stay above it.
#[test]
fn test_finer_sort_stays_above_merge_without_fetch() {
    let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
    let a = sort_expr_on("a", 0, false, false);
    let a_b = LexOrdering::new(
        a.iter()
            .cloned()
            .chain(sort_expr_on("b", 1, false, false).iter().cloned()),
    )
    .unwrap();
    let merge = Arc::new(SortPreservingMergeExec::new(a, source));
    let plan: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(a_b, merge));

    let mut sort_push_down = SortPushDown::new_default(plan);
    assign_initial_requirements(&mut sort_push_down);
    let pushed_down = plan_string(&pushdown_sorts(sort_push_down).unwrap().plan);

    assert_snapshot!(pushed_down, @r"
    SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false]
      SortPreservingMergeExec: [a@0 ASC NULLS LAST]
        MockMultiPartitionExec
    ");
}

@github-actions github-actions Bot added the core Core DataFusion crate label Nov 7, 2026
@galaxytemple

Copy link
Copy Markdown
Contributor Author

Thanks for the review and suggestions @jayzhan211 ! I've updated the PR to decline pushdown for SPM(both with and without fetch), and applied the test case updates.

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @galaxytemple !

@jayzhan211
jayzhan211 enabled auto-merge November 8, 2026 01:03
@jayzhan211
jayzhan211 added this pull request to the merge queue Nov 8, 2026
Merged via the queue into apache:main with commit 3f4a37f Nov 8, 2026
44 checks passed
alamb pushed a commit that referenced this pull request Nov 11, 2026
…27331)

## Which issue does this PR close?

- Backport of #27060 to
`branch-56`.
- Fixes #26162 on the 56
release line.
- Part of #24461.

## Rationale for this change

A finer ordering pushed below `SortPreservingMergeExec` is lost when the
merge keeps fewer keys. This produces incorrectly ordered results and
can make a window query fail plan validation. Downstream testing still
reproduces the failure in the 56 line.

The original fix merged after `branch-56` was cut and is needed before
the DataFusion 56 release candidate.

## What changes are included in this PR?

Cherry-pick `3f4a37fbabdcb07705db6341268e6c0b34da0d34` from #27060. Keep
the finer parent sort above the merge, together with the original
regression tests.

The `limit.slt` conflict was append context from an unrelated test
absent on `branch-56`; only the original #26162 regression block is
added.

## What is the testing strategy for this PR?

Carries the original `limit.slt` checks for ordered output, a window, an
ordered aggregate and a UNION, plus the physical-plan regression for a
merge without a fetch.

`cargo fmt --all`, both targeted physical-plan tests and the `limit.slt`
checks passed. Reverting the optimizer change makes the new
non-fetched-merge regression fail.

## Are there any user-facing changes?

Queries retain their requested ordering above sort-preserving merges. No
public API changes.

This PR was prepared with AI assistance.

Co-authored-by: Sanghoon <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) v57.0.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants