Repository navigation
fix: decline sort pushdown through SortPreservingMergeExec with fetch (#26162) - #27060
Conversation
…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.
|
While working on this, I also noticed that SPM without limit (at L589 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
left a comment
There was a problem hiding this comment.
@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() { | |||
There was a problem hiding this comment.
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
");
}Co-authored-by: Jay Zhan <[email protected]>
|
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
left a comment
There was a problem hiding this comment.
Thanks @galaxytemple !
…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]>
Which issue does this PR close?
Rationale for this change
In #26162, running
ORDER BY b, aover a multi-partition subquery withORDER BY b LIMIT nproduces incorrect ordering ona.Pushing
[b, a]down removed the outerSortExec, but SPM only merges on[b]. As a result,awas 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:pushdown_requirement_to_childrenonly adds requirements to children, I think a logic to rewrite SPM adds extra complexity.What changes are included in this PR?
Ok(None)frompushdown_requirement_to_childrenwhenplanisSortPreservingMergeExecand has afetch.limit.slt.What is the testing strategy for this PR?
limit.sltreproducing 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.Are there any user-facing changes?
Yes, bug fix: queries like
ORDER BY b, aover a multi-partition subquery withORDER BY b LIMIT nnow return correctly ordered results ona. No API changes.