[branch-56] fix: retain finer sorts above merges (backport #27060)

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<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
+    ");
+}
+
 // ========================================================================
 // Multiple sort columns test
 // ========================================================================

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::<SortPreservingMergeExec>() && 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::<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_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::<HashJoinExec>() {
         handle_hash_join(hash_join, parent_required)
     } else {

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;