jayzhan211 commented on code in PR #27060:
URL: https://github.com/apache/datafusion/pull/27060#discussion_r4499820634
##########
datafusion/sqllogictest/test_files/limit.slt:
##########
@@ -1331,3 +1331,43 @@ SET datafusion.optimizer.max_passes = 3;
statement ok
DROP TABLE t26431;
+
+# 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
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+CREATE TABLE t_26162 AS
+ SELECT 1 AS a, 1 AS b UNION ALL SELECT 5, 2 UNION ALL SELECT 9, 3
+ UNION ALL SELECT 2, 1 UNION ALL SELECT 6, 2 UNION ALL SELECT 10, 3
+ UNION ALL SELECT 3, 1 UNION ALL SELECT 7, 2 UNION ALL SELECT 11, 3
+ UNION ALL SELECT 4, 1 UNION ALL SELECT 8, 2 UNION ALL SELECT 12, 3;
Review Comment:
Optional cleanup: this query also passes without the fix, because `a`
already comes out sorted from these one-row partitions; only the `EXPLAIN`
below fails. Reversing `a` across partitions makes the query itself fail
without the fix (it returns `4 3 2 1 8 7 6 5`). The expected output and plan
stay the same.
```suggestion
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;
```
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs:
##########
@@ -525,20 +525,13 @@ fn pushdown_requirement_to_children(
Ok(None)
}
} else if plan.is::<SortPreservingMergeExec>() && plan.fetch().is_some() {
Review 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.
```diff
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>() {
```
<details>
<summary>Test that fails without the change</summary>
For `datafusion/core/tests/physical_optimizer/ensure_requirements.rs`:
```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
");
}
```
</details>
##########
datafusion/sqllogictest/test_files/limit.slt:
##########
@@ -1331,3 +1331,43 @@ SET datafusion.optimizer.max_passes = 3;
statement ok
DROP TABLE t26431;
+
+# 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
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+CREATE TABLE t_26162 AS
+ SELECT 1 AS a, 1 AS b UNION ALL SELECT 5, 2 UNION ALL SELECT 9, 3
+ UNION ALL SELECT 2, 1 UNION ALL SELECT 6, 2 UNION ALL SELECT 10, 3
+ UNION ALL SELECT 3, 1 UNION ALL SELECT 7, 2 UNION ALL SELECT 11, 3
+ UNION ALL SELECT 4, 1 UNION ALL SELECT 8, 2 UNION ALL SELECT 12, 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]
+
+statement ok
Review Comment:
Optional cleanup: this fix also turns three planning errors into correct
results, and none of them is tested here. Before this PR, a window, an ordered
aggregate or a UNION above the fetched merge fails with `SanityCheckPlan: ...
does not satisfy the order and distribution requirements`. These queries fail
that way without the fix and pass with it.
```suggestion
# 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
```
##########
datafusion/sqllogictest/test_files/limit.slt:
##########
@@ -1331,3 +1331,43 @@ SET datafusion.optimizer.max_passes = 3;
statement ok
DROP TABLE t26431;
+
+# 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
+set datafusion.execution.target_partitions = 4;
+
Review Comment:
Optional cleanup: this `SET` has no effect, since line 1198 already sets
`target_partitions` to 4 and nothing changes it afterwards.
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
