From 59e3754f96dc2da2c40d9c52e552e040a005ded6 Mon Sep 17 00:00:00 2001 From: Sanghoon <75651582+galaxytemple@users.noreply.github.com> Date: Sun, 8 Nov 2026 01:04:29 +0000 Subject: [PATCH] fix: decline sort pushdown through SortPreservingMergeExec with fetch (#26162) (#27060) ## Which issue does this PR close? - Closes #26162. ## 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? - query on a partitioned table in `limit.slt` reproducing #26162. ## 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. --------- Co-authored-by: Jay Zhan (cherry picked from commit 3f4a37fbabdcb07705db6341268e6c0b34da0d34) --- .../physical_optimizer/ensure_requirements.rs | 26 +++++++ .../enforce_sorting/sort_pushdown.rs | 39 ++-------- datafusion/sqllogictest/test_files/limit.slt | 71 +++++++++++++++++++ 3 files changed, 102 insertions(+), 34 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 86adb307db75b..0d779a74308cc 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -522,6 +522,32 @@ fn test_fetched_merge_keeps_stricter_topk_above() { "); } +/// 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 = 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 = 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 + "); +} + // ======================================================================== // Multiple sort columns test // ======================================================================== diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index a821c0e49c724..7882de3ea55fc 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -5,10 +5,7 @@ use std::fmt::Debug; use std::sync::Arc; -use crate::utils::{ - add_sort_above_with_distribution, is_sort, is_sort_preserving_merge, is_union, - is_window, -}; +use crate::utils::{add_sort_above_with_distribution, is_sort, is_union, is_window}; use arrow::datatypes::SchemaRef; use datafusion_common::tree_node::{Transformed, TreeNode}; @@ -524,21 +521,10 @@ fn pushdown_requirement_to_children( } else { Ok(None) } - } else if plan.is::() && plan.fetch().is_some() { - // Push down through a fetched merge when: - // - requirement is aligned with output ordering - // - it preserves ordering during execution - // Refining a fetched merge must preserve the ordering that selected its rows. - let Some(ordering) = plan.output_ordering() else { - return Ok(None); - }; - Ok(plan - .equivalence_properties() - .requirements_compatible( - parent_required.first().clone(), - ordering.clone().into(), - ) - .then(|| vec![Some(parent_required)])) + } else if plan.is::() { + // 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_union(plan) { // `UnionExec` does not have real sort requirements for its input, we // just propagate the sort requirements down: @@ -593,21 +579,6 @@ fn pushdown_requirement_to_children( // Keep sorts above RepartitionExec even when a single input partition // lets it preserve order; pushing the sort below it provides no benefit. 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::() { handle_hash_join(hash_join, parent_required) } else { diff --git a/datafusion/sqllogictest/test_files/limit.slt b/datafusion/sqllogictest/test_files/limit.slt index 6d6cea4a83f51..b26dfd1cc6d89 100644 --- a/datafusion/sqllogictest/test_files/limit.slt +++ b/datafusion/sqllogictest/test_files/limit.slt @@ -1299,3 +1299,74 @@ SET datafusion.optimizer.max_passes = 3; statement ok DROP TABLE t23534; + +# Test for #26162: ORDER BY b, a over a subquery with ORDER BY b LIMIT n +# Ensure SortPreservingMergeExec with fetch does not push down finer sort requirements +statement ok +CREATE TABLE t_26162 AS + SELECT 4 AS a, 1 AS b UNION ALL SELECT 8, 2 UNION ALL SELECT 12, 3 + UNION ALL SELECT 3, 1 UNION ALL SELECT 7, 2 UNION ALL SELECT 11, 3 + UNION ALL SELECT 2, 1 UNION ALL SELECT 6, 2 UNION ALL SELECT 10, 3 + UNION ALL SELECT 1, 1 UNION ALL SELECT 5, 2 UNION ALL SELECT 9, 3; + +query II +SELECT * FROM (SELECT a, b FROM t_26162 ORDER BY b LIMIT 8) ORDER BY b, a; +---- +1 1 +2 1 +3 1 +4 1 +5 2 +6 2 +7 2 +8 2 + +query TT +EXPLAIN SELECT * FROM (SELECT a, b FROM t_26162 ORDER BY b LIMIT 8) ORDER BY b, a; +---- +logical_plan +01)Sort: t_26162.b ASC NULLS LAST, t_26162.a ASC NULLS LAST +02)--Sort: t_26162.b ASC NULLS LAST, fetch=8 +03)----TableScan: t_26162 projection=[a, b] +physical_plan +01)SortExec: expr=[b@1 ASC NULLS LAST, a@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--SortPreservingMergeExec: [b@1 ASC NULLS LAST], fetch=8 +03)----SortExec: TopK(fetch=8), expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: partitions=12, partition_sizes=[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] + +# Requirements from a window, an ordered aggregate and a UNION also reach the merge. +query III +SELECT a, b, row_number() OVER (ORDER BY b, a) FROM (SELECT a, b FROM t_26162 ORDER BY b LIMIT 8) ORDER BY b, a; +---- +1 1 1 +2 1 2 +3 1 3 +4 1 4 +5 2 5 +6 2 6 +7 2 7 +8 2 8 + +query ? +SELECT array_agg(a ORDER BY b, a) FROM (SELECT a, b FROM t_26162 ORDER BY b LIMIT 8); +---- +[1, 2, 3, 4, 5, 6, 7, 8] + +query II +SELECT * FROM ( + SELECT * FROM (SELECT a, b FROM t_26162 ORDER BY b LIMIT 4) + UNION ALL + SELECT * FROM (SELECT a, b FROM t_26162 ORDER BY b LIMIT 4) +) ORDER BY b, a; +---- +1 1 +1 1 +2 1 +2 1 +3 1 +3 1 +4 1 +4 1 + +statement ok +DROP TABLE t_26162;