Skip to content

Commit ebf4858

Browse files
HarshilMitalHarshilMital
authored
fix: SortPreservingMergeExec drops input orderings when merging partitions (#26654)
## Which issue does this PR close? - Closes #26646. ## Rationale for this change When a sorted input with more than one partition goes through `SortPreservingMergeExec`, the planner treats every ordering of the input as still holding after the merge. Only the merge key ordering actually holds. The other orderings are true within each partition, but not across them. Because of this, the planner can remove a sort that is still needed, and the query returns rows out of order. With two files that are each sorted by `a` and `b` and `target_partitions = 2`: ```sql SELECT * FROM (SELECT a, b FROM t ORDER BY a LIMIT 10) ORDER BY b; ``` the plan has no sort for the outer `ORDER BY b`, and the rows come back as `(1, 10), (1, 1), (2, 20), (3, 2)`. ## What changes are included in this PR? - `SortPreservingMergeExec::compute_properties` clears the input orderings when there is more than one input partition, before adding the merge key ordering. This matches what `CoalescePartitionsExec` and `RepartitionExec` already do. With a single input partition nothing is merged, so the input orderings are kept as before. The preserve-order `RepartitionExec` case mentioned in the issue is not covered here. Draft #26643 moves this logic into `EquivalenceProperties::merged`. If that lands first, I'll update this PR to use `.merged(input_partitions, false)` followed by `add_ordering`. ## What is the testing strategy for this PR? - New sqllogictest case at the end of `order.slt`, using the reproduction from the issue. It checks both the query result and the physical plan, which must keep the `SortExec` above the `SortPreservingMergeExec`. It fails on `main` with the wrong row order from the issue, and passes with this change. - The existing sqllogictests and the `datafusion-physical-plan` / `datafusion-physical-optimizer` unit tests pass with no changes to expected plans. ## Are there any user-facing changes? Queries like the one above now return correctly ordered results. No public API changes.
1 parent cd17d8e commit ebf4858

2 files changed

Lines changed: 62 additions & 0 deletions

File tree

‎datafusion/physical-plan/src/sorts/sort_preserving_merge.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,11 @@ impl SortPreservingMergeExec {
172172
};
173173

174174
let mut eq_properties = input.equivalence_properties().clone();
175+
// Merging multiple partitions only preserves the merge key ordering.
176+
// Other input orderings hold only within each partition:
177+
if input_partitions > 1 {
178+
eq_properties.clear_orderings();
179+
}
175180
eq_properties.clear_per_partition_constants();
176181
eq_properties.add_ordering(ordering);
177182
PlanProperties::new(

‎datafusion/sqllogictest/test_files/order.slt‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2276,3 +2276,60 @@ SELECT i, -i AS n FROM (
22762276
-128 -128
22772277
-128 -128
22782278
NULL NULL
2279+
2280+
# SortPreservingMergeExec only guarantees its merge key ordering. Each file
2281+
# below is sorted by both `a` and `b`, but merging them on `a` does not keep
2282+
# them sorted by `b`, so the outer sort must stay.
2283+
statement ok
2284+
set datafusion.execution.target_partitions = 2;
2285+
2286+
query I
2287+
COPY (VALUES (1, 10), (2, 20))
2288+
TO 'test_files/scratch/order/spm_orderings/1.csv'
2289+
OPTIONS ('format.has_header' 'true');
2290+
----
2291+
2
2292+
2293+
query I
2294+
COPY (VALUES (1, 1), (3, 2))
2295+
TO 'test_files/scratch/order/spm_orderings/2.csv'
2296+
OPTIONS ('format.has_header' 'true');
2297+
----
2298+
2
2299+
2300+
statement ok
2301+
CREATE EXTERNAL TABLE spm_orderings (
2302+
a BIGINT,
2303+
b BIGINT
2304+
)
2305+
STORED AS CSV
2306+
WITH ORDER (a ASC)
2307+
WITH ORDER (b ASC)
2308+
LOCATION 'test_files/scratch/order/spm_orderings/'
2309+
OPTIONS ('format.has_header' 'true');
2310+
2311+
query II
2312+
SELECT * FROM (SELECT a, b FROM spm_orderings ORDER BY a LIMIT 10) ORDER BY b;
2313+
----
2314+
1 1
2315+
3 2
2316+
1 10
2317+
2 20
2318+
2319+
query TT
2320+
EXPLAIN SELECT * FROM (SELECT a, b FROM spm_orderings ORDER BY a LIMIT 10) ORDER BY b;
2321+
----
2322+
logical_plan
2323+
01)Sort: spm_orderings.b ASC NULLS LAST
2324+
02)--Sort: spm_orderings.a ASC NULLS LAST, fetch=10
2325+
03)----TableScan: spm_orderings projection=[a, b]
2326+
physical_plan
2327+
01)SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false]
2328+
02)--SortPreservingMergeExec: [a@0 ASC NULLS LAST], fetch=10
2329+
03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/order/spm_orderings/1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/order/spm_orderings/2.csv]]}, projection=[a, b], limit=10, output_orderings=[[a@0 ASC NULLS LAST], [b@1 ASC NULLS LAST]], file_type=csv, has_header=true
2330+
2331+
statement ok
2332+
DROP TABLE spm_orderings;
2333+
2334+
statement ok
2335+
set datafusion.execution.target_partitions = 4;

0 commit comments

Comments
 (0)