You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/23 11:09:57 UTC
(datafusion) branch gh-readonly-queue/main/pr-21976-258c6c2cd9113a10cc755b1c246134a31bf288bb created (now fcf9c9e43d)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-21976-258c6c2cd9113a10cc755b1c246134a31bf288bb
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at fcf9c9e43d Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts (#21976)
This branch includes the following new commits:
new fcf9c9e43d Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts (#21976)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts (#21976)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-21976-258c6c2cd9113a10cc755b1c246134a31bf288bb
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit fcf9c9e43db50693de6edc58325f3c83488a934f
Author: Qi Zhu <qi...@polygon.io>
AuthorDate: Sat May 23 19:09:34 2026 +0800
Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts (#21976)
## Summary
Replace the separate `EnforceDistribution` and `EnforceSorting`
optimizer rules with a single `EnsureRequirements` rule in the default
optimizer chain. Fix `pushdown_sorts` to be distribution-aware and fix
the `SortPreservingMergeExec` / `CoalescePartitionsExec` fetch
preservation issue from #14150, making the composition idempotent.
**Epic**: #21973
Closes: #14150
## Problem
`EnforceDistribution` and `EnforceSorting` run as separate rules, but
sorting and distribution are coupled through
`SortExec.preserve_partitioning`. This caused:
1. **`SanityCheckPlan` validation failures on multi-partition sort +
limit** — `pushdown_sorts` set `preserve_partitioning=true` on
multi-partition input without inserting `SortPreservingMergeExec`,
violating the `SinglePartition` requirement coming from
`GlobalLimitExec`.
2. **Non-idempotent composition** — running the rules multiple times
produced different (sometimes invalid) plans.
3. **Lost fetch values (#14150)** — `EnforceDistribution` dropped
`fetch` from `SortPreservingMergeExec` / `CoalescePartitionsExec` when
stripping and re-adding distribution operators.
DataFusion was the only major query engine with separate rules — Spark
(`EnsureRequirements`) and Presto/Trino (`AddExchanges`) handle both in
a single rule.
## Changes
### 1. `EnsureRequirements` rule (new, replaces `EnforceDistribution` +
`EnforceSorting` in the default chain)
- Single `PhysicalOptimizerRule` that calls the distribution + sorting
helpers in one coordinated bottom-up sequence.
- Registered in place of `Arc::new(EnforceDistribution) +
Arc::new(EnforceSorting)` in the default optimizer chain.
- Comprehensive inline tests covering known bug topologies + idempotency
verification.
### 2. Distribution-aware `pushdown_sorts` (`sort_pushdown.rs`)
- Add `distribution_requirement: Distribution` field to
`ParentRequirements`.
- New `add_sort_above_with_distribution()` in `utils.rs` — inserts
`SortPreservingMergeExec` when the parent requires `SinglePartition` and
the input has multiple partitions.
- Switch both `add_sort_above` call sites to the distribution-aware
variant.
- Propagate distribution through recursion with a
`stronger_distribution()` helper.
- Reset distribution below partition-merging nodes (SPM,
single-partition outputs).
### 3. Fix fetch preservation in distribution enforcement (#14150)
- `remove_dist_changing_operators()` now saves `fetch` from removed SPM
/ Coalesce nodes.
- `add_merge_on_top()` re-applies the saved `fetch` to re-created
operators.
### 4. Retire the old rule entry points; retarget existing tests
After review feedback from @alamb
(https://github.com/apache/datafusion/pull/21976#issuecomment-4432397769),
the rule structs and their `impl PhysicalOptimizerRule` blocks have been
deleted from `enforce_distribution.rs` and `enforce_sorting/mod.rs`. The
internal helpers (`ensure_distribution`, `ensure_sorting`, the contexts,
`parallelize_sorts`, `replace_with_order_preserving_variants`,
`sort_pushdown`, …) stay in place — `EnsureRequirements` calls them
directly.
The existing integration tests in `core/tests/physical_optimizer/` now
exercise `EnsureRequirements` instead of the deleted rules:
- `enforce_distribution.rs` — `Run::Distribution` / `Run::Sorting`
branches both call `EnsureRequirements::new()`. Legacy run sequences
(`DISTRIB_DISTRIB_SORT`, `SORT_DISTRIB_DISTRIB`) are preserved verbatim;
idempotency makes the previously-different orderings converge to the
same plan.
- `enforce_sorting.rs` — `EnforceSortingTest` drives
`EnsureRequirements::new()` and pins `target_partitions = 10` so
snapshots are deterministic across machines. The historical `[Dist,
Sort]` vs `[Sort, Dist, Sort]` comparison is rewritten as "running
`EnsureRequirements` N times == running it once".
- `enforce_sorting_monotonicity.rs` /
`replace_with_order_preserving_variants.rs` — driven through the same
test framework; only snapshots updated.
A previously-separate `ensure_requirements/new_tests.rs` (added in an
earlier iteration of this PR) is removed; the same coverage lives in the
inline tests in `ensure_requirements/mod.rs`.
### 5. Updated SLT
- `explain.slt`: `EnforceDistribution` + `EnforceSorting` collapse to
`EnsureRequirements` in `EXPLAIN VERBOSE` output.
## Snapshot drift
~78 snapshots in the retargeted tests refreshed. The consistent pattern
is `SortExec + CoalescePartitionsExec` (blocking) →
`SortPreservingMergeExec` (streaming), because `EnsureRequirements` now
runs `parallelize_sorts` + `replace_with_order_preserving_variants` on
plan shapes that the single-rule path used to miss. These are
improvements, not regressions — but worth a careful look in review since
they are visible in the diff.
## Testing
| Suite | Result |
|-------|--------|
| `datafusion-physical-optimizer` (lib, inline tests) | **59 passed** |
| `core_integration physical_optimizer::` | **454 passed** |
| `cargo clippy --all-targets -- -D warnings` | clean |
| `cargo fmt --all --check` | clean |
### Idempotency / regression coverage in the inline tests
| Scenario | Covered |
|----------|---------|
| Multi-partition sort + limit (1-64 partitions) | yes |
| Union with mixed partition counts | yes |
| Projection over multi-partition | yes |
| HashJoin (Partitioned) | yes |
| SortMergeJoin | yes |
| Window function partitioning + ordering | yes |
| Aggregate (Partial + FinalPartitioned) | yes |
| Nested sort + limit | yes |
| Hash repartition + sort | yes |
| CoalescePartitions + sort (`parallelize_sorts`) | yes |
| SPM → Sort → multi-partition | yes |
| `OutputRequirementExec` + `SinglePartition` over multi-partition
source | yes |
| `ProjectionExec` + multi-partition + `SinglePartition` requirement |
yes |
| #14150 fetch preservation across passes | yes |
| Triple optimization convergence | yes |
| 10× consecutive optimization stability | yes |
## Architecture
```
EnsureRequirements::optimize(plan)
Phase 1: join key reordering (top-down) — adjust_input_keys_ordering
or reorder_join_keys_to_inputs depending on config.
Phase 2: distribution enforcement (bottom-up) — ensure_distribution
Fetch is preserved across SPM/Coalesce strip/re-add (#14150 fix).
Phase 3: sort enforcement (bottom-up) — ensure_sorting
Phase 4: parallelize_sorts (bottom-up, when repartition_sorts is on)
Phase 5: replace_with_order_preserving_variants (bottom-up)
Phase 6: pushdown_sorts (top-down, distribution-aware)
Phase 7: replace_with_partial_sort (bottom-up)
```
Idempotent because:
- `pushdown_sorts` now carries `distribution_requirement` and uses
`add_sort_above_with_distribution`, so the second pass never re-violates
an earlier-established `SinglePartition` requirement.
- Distribution enforcement preserves `fetch` across strip/re-add cycles.
- Running `EnsureRequirements` repeatedly converges (verified across the
partition-count sweep, hash-join, sort-merge join, window, projection,
and #14150 regression tests).
## Next steps (future PRs)
- Gradually fold `pushdown_sorts` work into the bottom-up
`ensure_sorting` pass.
- Eliminate the separate top-down `pushdown_sorts` traversal.
- Single-pass architecture (one `transform_up` for both distribution +
sorting, like Spark's `EnsureRequirements`).
---
datafusion/core/src/optimizer_rule_reference.md | 31 +-
.../physical_optimizer/enforce_distribution.rs | 228 ++--
.../tests/physical_optimizer/enforce_sorting.rs | 361 +++---
.../enforce_sorting_monotonicity.rs | 261 +++-
.../physical_optimizer/ensure_requirements.rs | 1241 ++++++++++++++++++++
datafusion/core/tests/physical_optimizer/mod.rs | 1 +
.../src/combine_partial_final_agg.rs | 3 +-
.../enforce_distribution.rs | 258 ++--
.../enforce_sorting/mod.rs | 123 +-
.../replace_with_order_preserving_variants.rs | 0
.../enforce_sorting/sort_pushdown.rs | 125 +-
.../src/ensure_requirements/mod.rs | 259 ++++
datafusion/physical-optimizer/src/lib.rs | 7 +-
datafusion/physical-optimizer/src/optimizer.rs | 46 +-
datafusion/physical-optimizer/src/utils.rs | 52 +-
datafusion/sqllogictest/test_files/explain.slt | 12 +-
16 files changed, 2307 insertions(+), 701 deletions(-)
diff --git a/datafusion/core/src/optimizer_rule_reference.md b/datafusion/core/src/optimizer_rule_reference.md
index 1f9f37f530..7652c2dcae 100644
--- a/datafusion/core/src/optimizer_rule_reference.md
+++ b/datafusion/core/src/optimizer_rule_reference.md
@@ -75,20 +75,19 @@ in multiple phases.
| 3 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. |
| 4 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. |
| 5 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. |
-| 6 | `EnforceDistribution` | - | Adds repartitioning only where needed to satisfy physical distribution requirements. |
+| 6 | `EnsureRequirements` | - | Enforces both distribution and sorting requirements in a single idempotent rule. |
| 7 | `CombinePartialFinalAggregate` | - | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant. |
-| 8 | `EnforceSorting` | - | Adds or removes local sorts to satisfy required input orderings. |
-| 9 | `OptimizeAggregateOrder` | - | Updates aggregate expressions to use the best ordering once sort requirements are known. |
-| 10 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. |
-| 11 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. |
-| 12 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. |
-| 13 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. |
-| 14 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. |
-| 15 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. |
-| 16 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. |
-| 17 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. |
-| 18 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. |
-| 19 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. |
-| 20 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. |
-| 21 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. |
-| 22 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. |
+| 8 | `OptimizeAggregateOrder` | - | Updates aggregate expressions to use the best ordering once sort requirements are known. |
+| 9 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. |
+| 10 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. |
+| 11 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. |
+| 12 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. |
+| 13 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. |
+| 14 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. |
+| 15 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. |
+| 16 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. |
+| 17 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. |
+| 18 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. |
+| 19 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. |
+| 20 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. |
+| 21 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. |
diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
index 12abf79041..fb11657107 100644
--- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
+++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs
@@ -53,7 +53,7 @@ use datafusion_physical_expr_common::sort_expr::{
};
use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion_physical_optimizer::enforce_distribution::*;
-use datafusion_physical_optimizer::enforce_sorting::EnforceSorting;
+use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
use datafusion_physical_optimizer::output_requirements::OutputRequirements;
use datafusion_physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
@@ -672,17 +672,14 @@ impl TestConfig {
// TODO: End state payloads will be checked here.
}
- for run in optimizers_to_run {
- optimized = match run {
- Run::Distribution => {
- let optimizer = EnforceDistribution::new();
- optimizer.optimize(optimized, &self.config)?
- }
- Run::Sorting => {
- let optimizer = EnforceSorting::new();
- optimizer.optimize(optimized, &self.config)?
- }
- };
+ // With `EnsureRequirements`, distribution and sorting enforcement are
+ // composed into a single idempotent pass, so the historical sequence
+ // of `Run::Distribution` / `Run::Sorting` collapses to repeated calls
+ // of the same rule. The sequences are preserved so existing test
+ // assertions (which encode legacy run orders) remain stable.
+ for _run in optimizers_to_run {
+ let optimizer = EnsureRequirements::new();
+ optimized = optimizer.optimize(optimized, &self.config)?;
}
// Remove the ancillary output requirements operator when done:
@@ -1574,15 +1571,15 @@ fn multi_smj_joins() -> Result<()> {
assert_plan!(plan_distrib, @r"
SortMergeJoinExec: join_type=..., on=[(a@0, c@2)]
SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[c@2 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
");
}
@@ -1598,18 +1595,18 @@ fn multi_smj_joins() -> Result<()> {
_ => {
assert_plan!(plan_distrib, @r"
SortMergeJoinExec: join_type=..., on=[(a@0, c@2)]
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[c@2 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
");
}
@@ -1649,17 +1646,16 @@ fn multi_smj_joins() -> Result<()> {
// TODO(wiedld): show different test result if enforce distribution first.
assert_plan!(plan_sort, @r"
SortMergeJoinExec: join_type=..., on=[(a@0, c@2)]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
+ SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
+ RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
+ ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
- ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true
SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
@@ -1686,15 +1682,15 @@ fn multi_smj_joins() -> Result<()> {
assert_plan!(plan_distrib, @r"
SortMergeJoinExec: join_type=..., on=[(b1@6, c@2)]
SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[c@2 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
");
}
@@ -1703,18 +1699,18 @@ fn multi_smj_joins() -> Result<()> {
// TODO(wiedld): show different test result if enforce sorting first.
assert_plan!(plan_distrib, @r"
SortMergeJoinExec: join_type=..., on=[(b1@6, c@2)]
- SortExec: expr=[b1@6 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10
+ RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@6 ASC
+ SortExec: expr=[b1@6 ASC], preserve_partitioning=[true]
SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[b1@1 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- SortExec: expr=[c@2 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1
+ RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
");
}
@@ -1748,17 +1744,16 @@ fn multi_smj_joins() -> Result<()> {
// TODO(wiedld): show different test result if enforce distribution first.
assert_plan!(plan_sort, @r"
SortMergeJoinExec: join_type=..., on=[(b1@6, c@2)]
- RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[b1@6 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+ RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@6 ASC
+ SortExec: expr=[b1@6 ASC], preserve_partitioning=[true]
+ SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
+ RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
+ SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
+ ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[b1@1 ASC], preserve_partitioning=[false]
- ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true
SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
@@ -1852,25 +1847,21 @@ fn smj_join_key_ordering() -> Result<()> {
let plan_sort = test_config.to_plan(join, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort, @r"
SortMergeJoinExec: join_type=Inner, on=[(b3@1, b2@1), (a3@0, a2@0)]
- RepartitionExec: partitioning=Hash([b3@1, a3@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]
- ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]
- AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]
- RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10
- AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- RepartitionExec: partitioning=Hash([b2@1, a2@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- ProjectionExec: expr=[a@1 as a2, b@0 as b2]
- AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]
- RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10
- AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
+ SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[true]
+ ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]
+ ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]
+ AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]
+ RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10
+ AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
+ SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[true]
+ ProjectionExec: expr=[a@1 as a2, b@0 as b2]
+ AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]
+ RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10
+ AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
");
Ok(())
@@ -1914,9 +1905,8 @@ fn merge_does_not_need_sort() -> Result<()> {
let plan_sort = test_config.to_plan(exec, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
+ SortPreservingMergeExec: [a@0 ASC]
+ DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
");
Ok(())
@@ -2246,9 +2236,8 @@ fn repartition_ignores_sort_preserving_merge() -> Result<()> {
let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
+ SortPreservingMergeExec: [c@2 ASC]
+ DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
");
Ok(())
@@ -2286,11 +2275,10 @@ fn repartition_ignores_sort_preserving_merge_with_union() -> Result<()> {
let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- UnionExec
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
+ SortPreservingMergeExec: [c@2 ASC]
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
");
Ok(())
@@ -2412,8 +2400,8 @@ fn repartition_transitively_with_projection() -> Result<()> {
let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[sum@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
+ SortPreservingMergeExec: [sum@0 ASC]
+ SortExec: expr=[sum@0 ASC], preserve_partitioning=[true]
ProjectionExec: expr=[a@0 + b@1 as sum]
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
@@ -2524,8 +2512,8 @@ fn repartition_transitively_past_sort_with_filter() -> Result<()> {
let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
+ SortPreservingMergeExec: [a@0 ASC]
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
FilterExec: c@2 = 0
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
@@ -2575,8 +2563,8 @@ fn repartition_transitively_past_sort_with_projection_and_filter() -> Result<()>
let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
+ SortPreservingMergeExec: [a@0 ASC]
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c]
FilterExec: c@2 = 0
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
@@ -3094,11 +3082,10 @@ fn parallelization_sort_preserving_merge_with_union() -> Result<()> {
let plan_parquet_sort = test_config.to_plan(plan_parquet, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_parquet_sort,
@r"
- SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- UnionExec
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
+ SortPreservingMergeExec: [c@2 ASC]
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
");
// no SPM
// has coalesce
@@ -3115,11 +3102,10 @@ fn parallelization_sort_preserving_merge_with_union() -> Result<()> {
let plan_csv_sort = test_config.to_plan(plan_csv.clone(), &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_csv_sort,
@r"
- SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- UnionExec
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false
+ SortPreservingMergeExec: [c@2 ASC]
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false
");
// no SPM
// has coalesce
@@ -3451,8 +3437,8 @@ fn do_not_preserve_ordering_through_repartition() -> Result<()> {
let plan_sort = test_config.to_plan(physical_plan, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
+ SortPreservingMergeExec: [a@0 ASC]
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
FilterExec: c@2 = 0
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2
DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
@@ -3522,12 +3508,11 @@ fn do_not_preserve_ordering_through_repartition2() -> Result<()> {
let plan_sort = test_config.to_plan(physical_plan, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort,
@r"
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- FilterExec: c@2 = 0
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2
- DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
+ SortPreservingMergeExec: [a@0 ASC]
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
+ FilterExec: c@2 = 0
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2
+ DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet
");
Ok(())
@@ -3580,14 +3565,15 @@ fn do_not_put_sort_when_input_is_invalid() -> Result<()> {
config.execution.target_partitions = 10;
config.optimizer.enable_round_robin_repartition = true;
config.optimizer.prefer_existing_sort = false;
- let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?;
+ let dist_plan = EnsureRequirements::new().optimize(physical_plan, &config)?;
// Since at the start of the rule ordering requirement is not satisfied
// EnforceDistribution rule doesn't satisfy this requirement either.
assert_plan!(dist_plan, @r"
SortRequiredExec: [a@0 ASC]
- FilterExec: c@2 = 0
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
+ SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
+ FilterExec: c@2 = 0
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
");
Ok(())
@@ -3616,7 +3602,7 @@ fn put_sort_when_input_is_valid() -> Result<()> {
config.execution.target_partitions = 10;
config.optimizer.enable_round_robin_repartition = true;
config.optimizer.prefer_existing_sort = false;
- let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?;
+ let dist_plan = EnsureRequirements::new().optimize(physical_plan, &config)?;
// Since at the start of the rule ordering requirement is satisfied
// EnforceDistribution rule satisfy this requirement also.
assert_plan!(dist_plan, @r"
@@ -3769,8 +3755,8 @@ async fn test_distribute_sort_parquet() -> Result<()> {
test_config.to_plan(physical_plan.clone(), &[Run::Distribution]);
assert_plan!(plan_distribution,
@r"
- SortExec: expr=[c@2 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
+ SortPreservingMergeExec: [c@2 ASC]
+ SortExec: expr=[c@2 ASC], preserve_partitioning=[true]
DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet
");
diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs
index 40bcdbbd6e..9a459f2049 100644
--- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs
+++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs
@@ -51,10 +51,10 @@ use datafusion_physical_plan::sorts::sort::SortExec;
use datafusion_physical_plan::{displayable, get_plan_string, ExecutionPlan};
use datafusion::datasource::physical_plan::CsvSource;
use datafusion::datasource::listing::PartitionedFile;
-use datafusion_physical_optimizer::enforce_sorting::{EnforceSorting, PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, parallelize_sorts, ensure_sorting};
+use datafusion_physical_optimizer::enforce_sorting::{PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, parallelize_sorts, ensure_sorting};
use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext};
use datafusion_physical_optimizer::enforce_sorting::sort_pushdown::{SortPushDown, assign_initial_requirements, pushdown_sorts};
-use datafusion_physical_optimizer::enforce_distribution::EnforceDistribution;
+use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
use datafusion_physical_optimizer::output_requirements::OutputRequirementExec;
use datafusion_physical_optimizer::PhysicalOptimizerRule;
use datafusion::prelude::*;
@@ -117,11 +117,19 @@ impl EnforceSortingTest {
pub(crate) fn run(&self) -> String {
let mut config = ConfigOptions::new();
config.optimizer.repartition_sorts = self.repartition_sorts;
-
- // This file has 4 rules that use tree node, apply these rules as in the
- // EnforceSorting::optimize implementation
- // After these operations tree nodes should be in a consistent state.
- // This code block makes sure that these rules doesn't violate tree node integrity.
+ // Pin target_partitions so snapshots stay deterministic across
+ // machines with different CPU counts. Now that the underlying
+ // optimizer is `EnsureRequirements` (which performs distribution
+ // enforcement), the partition count appears in `Hash([…], N)`
+ // nodes in the output plan; without pinning, snapshots taken on
+ // an N-core machine fail on an M-core machine. 10 matches the
+ // existing convention in `enforce_distribution.rs`.
+ config.execution.target_partitions = 10;
+
+ // This file has 4 sub-rules that use tree node; apply them in the same
+ // order EnsureRequirements does internally. After these operations the
+ // tree nodes should be in a consistent state; this block exists to make
+ // sure those sub-rules don't violate tree node integrity.
{
let plan_requirements =
PlanWithCorrespondingSort::new_default(Arc::clone(&self.plan));
@@ -175,9 +183,9 @@ impl EnforceSortingTest {
let input_plan_string = displayable(self.plan.as_ref()).indent(true).to_string();
// Run the actual optimizer
- let optimized_physical_plan = EnforceSorting::new()
+ let optimized_physical_plan = EnsureRequirements::new()
.optimize(Arc::clone(&self.plan), &config)
- .expect("enforce_sorting failed");
+ .expect("ensure_requirements failed");
// Get string representation of the plan
let optimized_plan_string = displayable(optimized_physical_plan.as_ref())
@@ -218,9 +226,12 @@ async fn test_remove_unnecessary_sort5() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
Optimized Plan:
- HashJoinExec: mode=Partitioned, join_type=Inner, on=[(col_a@0, c@2)]
- DataSourceExec: partitions=1, partition_sizes=[0]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
+ SortPreservingMergeExec: [a@2 ASC]
+ HashJoinExec: mode=Partitioned, join_type=Inner, on=[(col_a@0, c@2)]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
+ DataSourceExec: partitions=1, partition_sizes=[0]
+ RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet
");
Ok(())
}
@@ -255,14 +266,11 @@ async fn test_do_not_remove_sort_with_limit() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC, non_nullable_col@1 ASC]
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2
- UnionExec
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
- LocalLimitExec: fetch=100
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
+ LocalLimitExec: fetch=100
+ SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
// We should keep the bottom `SortExec`.
Ok(())
@@ -282,12 +290,18 @@ async fn test_union_inputs_sorted() -> Result<()> {
let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true);
assert_snapshot!(test.run(), @r"
- Input / Optimized Plan:
+ Input Plan:
SortPreservingMergeExec: [nullable_col@0 ASC]
UnionExec
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+
+ Optimized Plan:
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
// should not add a sort at the output of the union, input plan should not be changed
@@ -313,12 +327,18 @@ async fn test_union_inputs_different_sorted() -> Result<()> {
let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true);
assert_snapshot!(test.run(), @r"
- Input / Optimized Plan:
+ Input Plan:
SortPreservingMergeExec: [nullable_col@0 ASC]
UnionExec
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet
SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+
+ Optimized Plan:
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
// should not add a sort at the output of the union, input plan should not be changed
@@ -353,12 +373,9 @@ async fn test_union_inputs_different_sorted2() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC, non_nullable_col@1 ASC]
- UnionExec
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
Ok(())
@@ -485,13 +502,11 @@ async fn test_union_inputs_different_sorted3() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC]
- UnionExec
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
// should adjust sorting in the first input of the union such that it is not unnecessarily fine
Ok(())
@@ -529,14 +544,12 @@ async fn test_union_inputs_different_sorted4() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC, non_nullable_col@1 ASC]
- UnionExec
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ UnionExec
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
Ok(())
@@ -583,12 +596,9 @@ async fn test_union_inputs_different_sorted5() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC]
- UnionExec
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
Ok(())
@@ -631,14 +641,10 @@ async fn test_union_inputs_different_sorted6() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC]
- UnionExec
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
// Should adjust the requirement in the third input of the union so
// that it is not unnecessarily fine.
@@ -664,13 +670,20 @@ async fn test_union_inputs_different_sorted7() -> Result<()> {
// Union has unnecessarily fine ordering below it. We should be able to replace them with absolutely necessary ordering.
let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true);
assert_snapshot!(test.run(), @r"
- Input / Optimized Plan:
+ Input Plan:
SortPreservingMergeExec: [nullable_col@0 ASC]
UnionExec
SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+
+ Optimized Plan:
+ UnionExec
+ SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
// Union preserves the inputs ordering, and we should not change any of the SortExecs under UnionExec
@@ -807,9 +820,10 @@ async fn test_soft_hard_requirements_remove_soft_requirement_without_pushdowns()
Optimized Plan:
ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as count]
- BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
+ BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
"#);
// TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed
// let expected_optimized = [
@@ -855,10 +869,11 @@ async fn test_soft_hard_requirements_remove_soft_requirement_without_pushdowns()
Optimized Plan:
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
- ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
- SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
"#);
// TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed
// let expected_optimized = [
@@ -918,10 +933,11 @@ async fn test_soft_hard_requirements_multiple_soft_requirements() -> Result<()>
Optimized Plan:
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
- ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
- SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
"#);
// TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed
// let expected_optimized = [
@@ -985,10 +1001,11 @@ async fn test_soft_hard_requirements_multiple_soft_requirements() -> Result<()>
Optimized Plan:
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
- ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
- SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
"#);
// TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed
// let expected_optimized = [
@@ -1054,10 +1071,11 @@ async fn test_soft_hard_requirements_multiple_sorts() -> Result<()> {
Optimized Plan:
SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
- ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
- SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col]
+ RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
"#);
// TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed
// let expected_optimized = [
@@ -1165,7 +1183,7 @@ async fn test_window_multi_path_sort() -> Result<()> {
// are not necessarily the same to be able to remove them.
let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true);
assert_snapshot!(test.run(), @r#"
- Input Plan:
+ Input / Optimized Plan:
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
SortPreservingMergeExec: [nullable_col@0 DESC NULLS LAST]
UnionExec
@@ -1173,13 +1191,6 @@ async fn test_window_multi_path_sort() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet
SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
-
- Optimized Plan:
- WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Range, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
- SortPreservingMergeExec: [nullable_col@0 ASC]
- UnionExec
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet
"#);
Ok(())
@@ -1270,14 +1281,12 @@ async fn test_union_inputs_different_sorted_with_limit() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC]
- UnionExec
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- GlobalLimitExec: skip=0, fetch=100
- LocalLimitExec: fetch=100
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 DESC NULLS LAST], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ UnionExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ GlobalLimitExec: skip=0, fetch=100
+ LocalLimitExec: fetch=100
+ SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 DESC NULLS LAST], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
");
Ok(())
@@ -1346,10 +1355,12 @@ async fn test_sort_merge_join_order_by_left() -> Result<()> {
Optimized Plan:
SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)]
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
");
}
_ => {
@@ -1362,11 +1373,12 @@ async fn test_sort_merge_join_order_by_left() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
Optimized Plan:
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
- SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)]
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)]
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false]
+ SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
");
}
@@ -1436,10 +1448,12 @@ async fn test_sort_merge_join_order_by_right() -> Result<()> {
Optimized Plan:
SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)]
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[col_a@0 ASC, col_b@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
");
}
JoinType::RightAnti => {
@@ -1453,10 +1467,12 @@ async fn test_sort_merge_join_order_by_right() -> Result<()> {
Optimized Plan:
SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)]
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[col_a@0 ASC, col_b@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
");
}
_ => {
@@ -1469,11 +1485,12 @@ async fn test_sort_merge_join_order_by_right() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
Optimized Plan:
- SortExec: expr=[col_a@2 ASC, col_b@3 ASC], preserve_partitioning=[false]
- SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)]
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)]
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false]
+ SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
");
}
@@ -1518,11 +1535,12 @@ async fn test_sort_merge_join_complex_order_by() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
Optimized Plan:
- SortExec: expr=[col_b@3 ASC, nullable_col@0 ASC], preserve_partitioning=[false]
- SortMergeJoinExec: join_type=Inner, on=[(nullable_col@0, col_a@0)]
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ SortMergeJoinExec: join_type=Inner, on=[(nullable_col@0, col_a@0)]
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false]
+ SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
");
// can not push down the sort requirements, need to add SortExec
@@ -1546,10 +1564,12 @@ async fn test_sort_merge_join_complex_order_by() -> Result<()> {
Optimized Plan:
SortMergeJoinExec: join_type=Inner, on=[(nullable_col@0, col_a@0)]
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
- SortExec: expr=[col_a@0 ASC, col_b@1 ASC], preserve_partitioning=[false]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet
+ SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true]
+ RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet
");
// Can push down the sort requirements since col_a = nullable_col
@@ -1628,10 +1648,7 @@ async fn test_with_lost_ordering_unbounded() -> Result<()> {
StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC]
Optimized Plan:
- SortPreservingMergeExec: [a@0 ASC]
- RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
- StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC]
+ StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC]
");
let test_with_repartition_sorts =
@@ -1646,10 +1663,7 @@ async fn test_with_lost_ordering_unbounded() -> Result<()> {
StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC]
Optimized Plan:
- SortPreservingMergeExec: [a@0 ASC]
- RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
- StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC]
+ StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC]
");
Ok(())
@@ -1663,12 +1677,15 @@ async fn test_with_lost_ordering_bounded() -> Result<()> {
EnforceSortingTest::new(physical_plan.clone()).with_repartition_sorts(false);
assert_snapshot!(test_no_repartition_sorts.run(), @r"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
CoalescePartitionsExec
RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false
+
+ Optimized Plan:
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false
");
let test_with_repartition_sorts =
@@ -1683,11 +1700,7 @@ async fn test_with_lost_ordering_bounded() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false
Optimized Plan:
- SortPreservingMergeExec: [a@0 ASC]
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false
");
Ok(())
@@ -1705,11 +1718,15 @@ async fn test_do_not_pushdown_through_spm() -> Result<()> {
let test =
EnforceSortingTest::new(physical_plan.clone()).with_repartition_sorts(true);
assert_snapshot!(test.run(), @r"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[b@1 ASC], preserve_partitioning=[false]
SortPreservingMergeExec: [a@0 ASC, b@1 ASC]
RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false
+
+ Optimized Plan:
+ SortExec: expr=[b@1 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false
");
Ok(())
@@ -1741,10 +1758,8 @@ async fn test_pushdown_through_spm() -> Result<()> {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false
Optimized Plan:
- SortPreservingMergeExec: [a@0 ASC, b@1 ASC]
- SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false
+ SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false
");
Ok(())
}
@@ -1773,11 +1788,8 @@ async fn test_window_multi_layer_requirement() -> Result<()> {
Optimized Plan:
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- SortPreservingMergeExec: [a@0 ASC, b@1 ASC]
- SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false
+ SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false
"#);
Ok(())
@@ -1900,8 +1912,7 @@ async fn test_add_required_sort() -> Result<()> {
DataSourceExec: partitions=1, partition_sizes=[0]
Optimized Plan:
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: partitions=1, partition_sizes=[0]
+ DataSourceExec: partitions=1, partition_sizes=[0]
");
Ok(())
@@ -1967,9 +1978,8 @@ async fn test_remove_unnecessary_sort2() -> Result<()> {
DataSourceExec: partitions=1, partition_sizes=[0]
Optimized Plan:
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: partitions=1, partition_sizes=[0]
+ SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
+ DataSourceExec: partitions=1, partition_sizes=[0]
");
Ok(())
@@ -2012,9 +2022,7 @@ async fn test_remove_unnecessary_sort3() -> Result<()> {
Optimized Plan:
AggregateExec: mode=Final, gby=[], aggr=[]
- CoalescePartitionsExec
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: partitions=1, partition_sizes=[0]
+ DataSourceExec: partitions=1, partition_sizes=[0]
");
Ok(())
@@ -2064,10 +2072,8 @@ async fn test_remove_unnecessary_sort4() -> Result<()> {
SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[true]
FilterExec: NOT non_nullable_col@1
UnionExec
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: partitions=1, partition_sizes=[0]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: partitions=1, partition_sizes=[0]
+ DataSourceExec: partitions=1, partition_sizes=[0]
+ DataSourceExec: partitions=1, partition_sizes=[0]
");
Ok(())
@@ -2215,8 +2221,7 @@ async fn test_remove_unnecessary_spm1() -> Result<()> {
DataSourceExec: partitions=1, partition_sizes=[0]
Optimized Plan:
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
- DataSourceExec: partitions=1, partition_sizes=[0]
+ DataSourceExec: partitions=1, partition_sizes=[0]
");
Ok(())
@@ -2239,9 +2244,7 @@ async fn test_remove_unnecessary_spm2() -> Result<()> {
DataSourceExec: partitions=1, partition_sizes=[0]
Optimized Plan:
- LocalLimitExec: fetch=100
- SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false]
- DataSourceExec: partitions=1, partition_sizes=[0]
+ DataSourceExec: partitions=1, partition_sizes=[0]
");
Ok(())
@@ -2267,7 +2270,7 @@ async fn test_change_wrong_sorting() -> Result<()> {
DataSourceExec: partitions=1, partition_sizes=[0]
Optimized Plan:
- SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false]
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[0]
");
@@ -2295,7 +2298,7 @@ async fn test_change_wrong_sorting2() -> Result<()> {
DataSourceExec: partitions=1, partition_sizes=[0]
Optimized Plan:
- SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false]
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
DataSourceExec: partitions=1, partition_sizes=[0]
");
@@ -2360,22 +2363,16 @@ async fn test_commutativity() -> Result<()> {
"#);
let config = ConfigOptions::new();
- let rules = vec![
- Arc::new(EnforceDistribution::new()) as Arc<dyn PhysicalOptimizerRule>,
- Arc::new(EnforceSorting::new()) as Arc<dyn PhysicalOptimizerRule>,
- ];
- let mut first_plan = orig_plan.clone();
- for rule in rules {
- first_plan = rule.optimize(first_plan, &config)?;
- }
-
- let rules = vec![
- Arc::new(EnforceSorting::new()) as Arc<dyn PhysicalOptimizerRule>,
- Arc::new(EnforceDistribution::new()) as Arc<dyn PhysicalOptimizerRule>,
- Arc::new(EnforceSorting::new()) as Arc<dyn PhysicalOptimizerRule>,
- ];
+ // Idempotency check: under the previous design this verified that
+ // `[EnforceDistribution, EnforceSorting]` produced the same plan as
+ // `[EnforceSorting, EnforceDistribution, EnforceSorting]`. With the
+ // merged `EnsureRequirements` rule the property collapses to
+ // "running EnsureRequirements N times is the same as running it once",
+ // which is the idempotency guarantee the merged rule provides.
+ let rule = EnsureRequirements::new();
+ let first_plan = rule.optimize(orig_plan.clone(), &config)?;
let mut second_plan = orig_plan.clone();
- for rule in rules {
+ for _ in 0..3 {
second_plan = rule.optimize(second_plan, &config)?;
}
@@ -2414,10 +2411,8 @@ async fn test_coalesce_propagate() -> Result<()> {
DataSourceExec: partitions=1, partition_sizes=[0]
Optimized Plan:
- SortPreservingMergeExec: [nullable_col@0 ASC]
- SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: partitions=1, partition_sizes=[0]
+ SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false]
+ DataSourceExec: partitions=1, partition_sizes=[0]
");
Ok(())
diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs
index de7611ff21..99e25a6c82 100644
--- a/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs
+++ b/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs
@@ -433,8 +433,10 @@ fn test_window_partial_constant_and_set_monotonicity_8() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST]
+ WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -457,8 +459,10 @@ fn test_window_partial_constant_and_set_monotonicity_9() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST]
+ WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -475,10 +479,17 @@ fn test_window_partial_constant_and_set_monotonicity_10() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -495,10 +506,17 @@ fn test_window_partial_constant_and_set_monotonicity_11() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -520,10 +538,17 @@ fn test_window_partial_constant_and_set_monotonicity_12() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [non_nullable_col@1 ASC NULLS LAST, count@2 ASC NULLS LAST]
+ SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -541,10 +566,17 @@ fn test_window_partial_constant_and_set_monotonicity_13() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [non_nullable_col@1 ASC NULLS LAST, max@2 DESC NULLS LAST]
+ SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -562,10 +594,17 @@ fn test_window_partial_constant_and_set_monotonicity_14() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[min@2 DESC NULLS LAST, non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [min@2 DESC NULLS LAST, non_nullable_col@1 ASC NULLS LAST]
+ SortExec: expr=[min@2 DESC NULLS LAST, non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -583,10 +622,17 @@ fn test_window_partial_constant_and_set_monotonicity_15() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -806,8 +852,10 @@ fn test_window_partial_constant_and_set_monotonicity_24() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 DESC NULLS LAST]
+ WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -825,10 +873,17 @@ fn test_window_partial_constant_and_set_monotonicity_25() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -845,10 +900,17 @@ fn test_window_partial_constant_and_set_monotonicity_26() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [min@2 DESC NULLS LAST]
+ SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#);
}
@@ -865,10 +927,17 @@ fn test_window_partial_constant_and_set_monotonicity_27() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[avg@2 DESC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [avg@2 DESC NULLS LAST]
+ SortExec: expr=[avg@2 DESC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#);
}
@@ -891,10 +960,17 @@ fn test_window_partial_constant_and_set_monotonicity_28() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[count@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [count@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[count@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -918,8 +994,10 @@ fn test_window_partial_constant_and_set_monotonicity_29() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 DESC]
+ WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#)
}
@@ -935,10 +1013,17 @@ fn test_window_partial_constant_and_set_monotonicity_30() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [min@2 DESC NULLS LAST]
+ SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#);
}
@@ -955,10 +1040,17 @@ fn test_window_partial_constant_and_set_monotonicity_31() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[false]
WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[true]
+ WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1179,8 +1271,10 @@ fn test_window_partial_constant_and_set_monotonicity_40() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST]
+ BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1198,10 +1292,17 @@ fn test_window_partial_constant_and_set_monotonicity_41() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[max@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [max@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[max@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1219,10 +1320,17 @@ fn test_window_partial_constant_and_set_monotonicity_42() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1240,10 +1348,17 @@ fn test_window_partial_constant_and_set_monotonicity_43() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1265,10 +1380,17 @@ fn test_window_partial_constant_and_set_monotonicity_44() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[count@2 ASC], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [count@2 ASC]
+ SortExec: expr=[count@2 ASC], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1286,10 +1408,17 @@ fn test_window_partial_constant_and_set_monotonicity_45() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1313,8 +1442,10 @@ fn test_window_partial_constant_and_set_monotonicity_46() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST]
+ BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1337,8 +1468,10 @@ fn test_window_partial_constant_and_set_monotonicity_47() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST]
+ BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1553,8 +1686,10 @@ fn test_window_partial_constant_and_set_monotonicity_56() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [count@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1572,10 +1707,17 @@ fn test_window_partial_constant_and_set_monotonicity_57() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1593,10 +1735,17 @@ fn test_window_partial_constant_and_set_monotonicity_58() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST]
+ SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1613,10 +1762,17 @@ fn test_window_partial_constant_and_set_monotonicity_59() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[avg@2 ASC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [avg@2 ASC NULLS LAST]
+ SortExec: expr=[avg@2 ASC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1639,10 +1795,17 @@ fn test_window_partial_constant_and_set_monotonicity_60() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1660,10 +1823,17 @@ fn test_window_partial_constant_and_set_monotonicity_61() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 ASC]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1681,10 +1851,17 @@ fn test_window_partial_constant_and_set_monotonicity_62() {
],
}.run(),
@ r#"
- Input / Optimized Plan:
+ Input Plan:
SortExec: expr=[nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST], preserve_partitioning=[false]
BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+
+ Optimized Plan:
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST]
+ SortExec: expr=[nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST], preserve_partitioning=[true]
+ BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
@@ -1707,8 +1884,10 @@ fn test_window_partial_constant_and_set_monotonicity_63() {
DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
Optimized Plan:
- BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
+ SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST]
+ BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted]
+ RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet
"#
);
}
diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs
new file mode 100644
index 0000000000..16d58134c0
--- /dev/null
+++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs
@@ -0,0 +1,1241 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Integration tests for `EnsureRequirements`.
+//!
+//! Ported verbatim from `datafusion/physical-optimizer/src/ensure_requirements/mod.rs`
+//! so the tests live alongside the rest of the `physical_optimizer/` integration
+//! suite and can use real `ExecutionPlan`s where convenient.
+
+use datafusion_common::config::ConfigOptions;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
+
+use std::sync::Arc;
+
+use arrow::compute::SortOptions;
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion_common::Result;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+};
+use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+use datafusion_physical_plan::limit::GlobalLimitExec;
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::union::UnionExec;
+use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning,
+ PlanProperties, SendableRecordBatchStream,
+};
+
+use datafusion_physical_optimizer::output_requirements::OutputRequirementExec;
+use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan;
+
+use datafusion_common::{JoinType, NullEquality};
+use datafusion_physical_expr::Distribution;
+use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec};
+use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::repartition::RepartitionExec;
+use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+/// Mock ExecutionPlan with configurable partition count and output ordering.
+#[derive(Debug)]
+struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+}
+
+impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+}
+
+impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+}
+
+impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+}
+
+/// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+}
+
+/// Helper: verify idempotency — running twice produces the same plan
+fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = plan_string(&p1);
+ let s2 = plan_string(&p2);
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+}
+
+/// Single-column `LexOrdering` on `(name, idx)` with the given options.
+/// Most tests in this file want a one-column ordering on the canonical
+/// `a@0` or `b@1` columns; this helper trims the 7-line per-test
+/// boilerplate down to a single call.
+fn sort_expr_on(
+ name: &str,
+ idx: usize,
+ descending: bool,
+ nulls_first: bool,
+) -> LexOrdering {
+ LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new(name, idx)),
+ SortOptions {
+ descending,
+ nulls_first,
+ },
+ )])
+ .unwrap()
+}
+
+/// Render an execution plan with `displayable(...).indent(true)`.
+fn plan_string(plan: &Arc<dyn ExecutionPlan>) -> String {
+ datafusion_physical_plan::displayable(plan.as_ref())
+ .indent(true)
+ .to_string()
+}
+
+/// Run `EnsureRequirements`, assert `SanityCheckPlan` passes, snapshot the
+/// resulting plan with `insta`, and verify idempotency by running the rule a
+/// second time and checking the plan is unchanged. Use this for plan-shape
+/// tests so a single call covers "correct plan + sanity + idempotent" and
+/// updating an intentional plan change is a single `cargo insta accept`.
+macro_rules! assert_ensure_requirements_plan {
+ ($plan:expr, @ $snapshot:literal $(,)?) => {{
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize($plan, &config)
+ .expect("EnsureRequirements::optimize failed (pass 1)");
+ SanityCheckPlan::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("SanityCheckPlan failed (pass 1)");
+ let p1_str = plan_string(&p1);
+ insta::assert_snapshot!(p1_str, @ $snapshot);
+
+ // Idempotency: a second pass must produce the same plan.
+ let p2 = EnsureRequirements::new()
+ .optimize(p1, &config)
+ .expect("EnsureRequirements::optimize failed (pass 2)");
+ let p2_str = plan_string(&p2);
+ assert_eq!(
+ p1_str, p2_str,
+ "EnsureRequirements is NOT idempotent!\nPass 1:\n{p1_str}\nPass 2:\n{p2_str}",
+ );
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed (pass 2)");
+ }};
+}
+
+/// Union with mixed partition counts + sort + limit.
+#[test]
+fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=21
+ SortPreservingMergeExec: [a@0 DESC]
+ UnionExec
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[true]
+ MockMultiPartitionExec
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[false]
+ MockMultiPartitionExec
+ ");
+}
+
+/// Idempotency: union with mixed partitions
+#[test]
+fn test_idempotent_union_mixed_partitions() {
+ let live = Arc::new(MockMultiPartitionExec::new(8));
+ let hist = Arc::new(MockMultiPartitionExec::new(1));
+ let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap();
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ assert_idempotent(limit);
+}
+
+// ========================================================================
+// Projection + multi-partition tests (pushdown_sorts trigger path)
+// ========================================================================
+
+/// ProjectionExec over multi-partition + sort DESC + limit.
+/// This is the topology where pushdown_sorts pushes sort through projection
+/// onto the multi-partition source. The optimizer must still produce a valid plan.
+#[test]
+fn test_projection_over_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection = Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap());
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=21
+ SortPreservingMergeExec: [a@0 DESC]
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[true]
+ ProjectionExec: expr=[a@0 as a, b@1 as b]
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// Single partition tests (no unnecessary operators)
+// ========================================================================
+
+/// Single partition source + sort + limit should NOT add SortPreservingMergeExec.
+#[test]
+fn test_single_partition_no_unnecessary_spm() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ // Snapshot asserts the plan-shape property: no `SortPreservingMergeExec`
+ // is added on a single-partition source.
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=10
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[false]
+ MockMultiPartitionExec
+ ");
+}
+
+/// Source already has correct ordering → should not add SortExec.
+#[test]
+fn test_sort_already_satisfied_no_extra_sort() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ // Sort ASC matches MockMultiPartitionExec's output ordering (a ASC)
+ let sort_expr = sort_expr_on("a", 0, false, false);
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ // Snapshot asserts the plan-shape property: no `SortExec` is added
+ // when the source already satisfies the ordering.
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=10
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// Various partition counts (stress test)
+// ========================================================================
+
+/// Test with different partition counts: 2, 4, 8, 16, 32, 64
+#[test]
+fn test_various_partition_counts_all_pass_sanity_check() {
+ for n in [2, 4, 8, 16, 32, 64] {
+ let source = Arc::new(MockMultiPartitionExec::new(n));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(
+ result.is_ok(),
+ "SanityCheckPlan failed for {n} partitions: {:?}",
+ result.err()
+ );
+ }
+}
+
+// ========================================================================
+// CoalescePartitionsExec tests
+// ========================================================================
+
+/// CoalescePartitionsExec + sort should produce valid plan
+#[test]
+fn test_coalesce_then_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(8));
+ let coalesce: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(source));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, coalesce));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=10
+ SortPreservingMergeExec: [a@0 DESC]
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[true]
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// Filter + multi-partition tests
+// ========================================================================
+
+/// FilterExec over multi-partition + sort + limit
+#[test]
+fn test_filter_over_multi_partition_sort_limit() {
+ use datafusion_common::ScalarValue;
+ use datafusion_physical_expr::expressions::Literal;
+ use datafusion_physical_plan::filter::FilterExec;
+
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ // Simple always-true filter
+ let predicate = Arc::new(Literal::new(ScalarValue::Boolean(Some(true))));
+ let filter = Arc::new(FilterExec::try_new(predicate, source as _).unwrap());
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, filter));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=21
+ SortPreservingMergeExec: [a@0 DESC]
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[true]
+ FilterExec: true
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// RepartitionExec tests
+// ========================================================================
+
+/// Existing RepartitionExec + sort + limit must remain valid
+#[test]
+fn test_repartition_sort_limit_idempotent() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+ let repartition = Arc::new(
+ RepartitionExec::try_new(source as _, Partitioning::RoundRobinBatch(8)).unwrap(),
+ );
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, repartition));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=10
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[false]
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// Skip + Fetch (offset + limit) tests
+// ========================================================================
+
+/// GlobalLimitExec with skip=5, fetch=10 must produce valid plan
+#[test]
+fn test_skip_and_fetch_multi_partition() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ // skip=5, fetch=10
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 5, Some(10)));
+
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=5, fetch=10
+ SortPreservingMergeExec: [a@0 DESC]
+ SortExec: expr=[a@0 DESC], preserve_partitioning=[true]
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// Multiple sort columns test
+// ========================================================================
+
+/// Sort on (a DESC, b ASC) with multi-partition
+#[test]
+fn test_multi_column_sort_multi_partition() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![
+ PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ ),
+ PhysicalSortExpr::new(
+ Arc::new(Column::new("b", 1)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ ),
+ ])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ assert_ensure_requirements_plan!(limit, @r"
+ GlobalLimitExec: skip=0, fetch=21
+ SortPreservingMergeExec: [a@0 DESC, b@1 ASC NULLS LAST]
+ SortExec: expr=[a@0 DESC, b@1 ASC NULLS LAST], preserve_partitioning=[true]
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// pushdown_sorts distribution-awareness regression tests.
+// These cover the specific bug where pushdown_sorts pushed a SortExec
+// through an intermediate node onto a multi-partition source, setting
+// preserve_partitioning=true without inserting SortPreservingMergeExec.
+// ========================================================================
+
+/// Regression: `OutputRequirementExec(SinglePartition)` wrapping a
+/// multi-partition source. `ensure_sorting` must insert a
+/// `SortPreservingMergeExec` to satisfy the `SinglePartition` requirement.
+///
+/// The final `ensure_distribution` pass catches the distribution
+/// violation from `pushdown_sorts`, producing a valid plan (via
+/// `CoalescePartitionsExec` or `SortPreservingMergeExec`).
+#[test]
+fn test_output_requirement_single_partition_over_multi_partition_source() {
+ let source = Arc::new(MockMultiPartitionExec::new(10));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ // OutputRequirementExec with SinglePartition + ordering requirement
+ let output_req: Arc<dyn ExecutionPlan> = Arc::new(OutputRequirementExec::new(
+ source,
+ Some(OrderingRequirements::from(sort_expr)),
+ Distribution::SinglePartition,
+ Some(21),
+ ));
+
+ // SinglePartition must be satisfied (via SPM or Coalesce+Sort) — snapshot
+ // documents which one the optimizer chooses.
+ assert_ensure_requirements_plan!(output_req, @r"
+ OutputRequirementExec: order_by=[(a@0, desc)], dist_by=SinglePartition
+ SortPreservingMergeExec: [a@0 DESC], fetch=21
+ SortExec: TopK(fetch=21), expr=[a@0 DESC], preserve_partitioning=[true]
+ MockMultiPartitionExec
+ ");
+}
+
+/// Regression: `pushdown_sorts` pushes a sort through a `ProjectionExec`
+/// onto a multi-partition source. The result must include a
+/// `SortPreservingMergeExec` (or equivalent) when the parent requires
+/// `SinglePartition`.
+///
+/// Without the distribution-aware pushdown the standalone
+/// `pushdown_sorts` traversal would not propagate distribution; the
+/// final `ensure_distribution` pass then catches the violation and
+/// inserts a `CoalescePartitionsExec` to satisfy `SinglePartition`.
+#[test]
+fn test_sort_pushdown_through_projection_adds_spm() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(10));
+
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection: Arc<dyn ExecutionPlan> =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source).unwrap());
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ // OutputRequirementExec(SinglePartition) → ProjectionExec → multi-partition source
+ let output_req: Arc<dyn ExecutionPlan> = Arc::new(OutputRequirementExec::new(
+ projection,
+ Some(OrderingRequirements::from(sort_expr)),
+ Distribution::SinglePartition,
+ Some(21),
+ ));
+
+ // SinglePartition must be satisfied. The final ensure_distribution pass
+ // adds CoalescePartitionsExec or SortPreservingMergeExec as needed — the
+ // snapshot documents which.
+ assert_ensure_requirements_plan!(output_req, @r"
+ OutputRequirementExec: order_by=[(a@0, desc)], dist_by=SinglePartition
+ SortPreservingMergeExec: [a@0 DESC], fetch=21
+ SortExec: TopK(fetch=21), expr=[a@0 DESC], preserve_partitioning=[true]
+ ProjectionExec: expr=[a@0 as a, b@1 as b]
+ MockMultiPartitionExec
+ ");
+}
+
+// ========================================================================
+// Idempotency tests for distribution-fix scenarios
+// These verify that the pushdown_sorts distribution fix actually
+// makes EnsureRequirements idempotent for the bug-triggering topologies.
+// ========================================================================
+
+/// Idempotency for the `OutputRequirementExec(SinglePartition)` +
+/// multi-partition source scenario. Running twice must produce the same plan.
+#[test]
+fn test_idempotent_output_requirement_single_partition() {
+ let source = Arc::new(MockMultiPartitionExec::new(10));
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let output_req: Arc<dyn ExecutionPlan> = Arc::new(OutputRequirementExec::new(
+ source,
+ Some(OrderingRequirements::from(sort_expr)),
+ Distribution::SinglePartition,
+ Some(21),
+ ));
+
+ assert_idempotent(output_req);
+}
+
+/// Idempotency for the `OutputRequirementExec(SinglePartition)` →
+/// `ProjectionExec` → multi-partition source scenario.
+#[test]
+fn test_idempotent_projection_over_multi_partition_with_single_partition_requirement() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(10));
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection: Arc<dyn ExecutionPlan> =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source).unwrap());
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let output_req: Arc<dyn ExecutionPlan> = Arc::new(OutputRequirementExec::new(
+ projection,
+ Some(OrderingRequirements::from(sort_expr)),
+ Distribution::SinglePartition,
+ Some(21),
+ ));
+
+ assert_idempotent(output_req);
+}
+
+/// Regression for #21973 (the issue this PR fixes).
+///
+/// Topology: `SPM → Sort(preserve=true) → multi-partition`. Under the old
+/// two-rule pipeline `EnforceSorting::pushdown_sorts` could mutate
+/// `preserve_partitioning` after `EnforceDistribution` had settled
+/// distribution, so pass 2 could regress this parallel plan into a serial
+/// one. This is the exact topology that caused
+/// [`test_pushdown_through_spm`](../enforce_sorting.rs) to fail before this
+/// PR; `EnsureRequirements` must keep the SPM-over-parallel-sort shape
+/// stable across passes.
+///
+/// Also acts as the "no extra SPM when already optimal" check — the
+/// input plan already contains exactly one `SortPreservingMergeExec`,
+/// so the optimised plan must too (we used to have a separate test for
+/// this property, but on this input it is implied by idempotency
+/// combined with the SPM-count assertion).
+#[test]
+fn test_issue_21973_idempotent_spm_sort_multi_partition() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(10));
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(
+ SortExec::new(sort_expr.clone(), source).with_preserve_partitioning(true),
+ );
+ let spm: Arc<dyn ExecutionPlan> =
+ Arc::new(SortPreservingMergeExec::new(sort_expr, sort));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(spm, 0, Some(21)));
+
+ // No-extra-SPM property: count SortPreservingMergeExec occurrences in
+ // the first optimisation pass — must be ≤ 1 (the original SPM survives,
+ // none are added).
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new()
+ .optimize(Arc::clone(&limit), &config)
+ .expect("optimize failed");
+ let plan_str = plan_string(&optimized);
+ let spm_count = plan_str.matches("SortPreservingMergeExec").count();
+ assert!(
+ spm_count <= 1,
+ "Extra SortPreservingMergeExec added ({spm_count} found):\n{plan_str}"
+ );
+
+ assert_idempotent(limit);
+}
+
+/// Regression for #21973: the `parallelize_sorts` rewrite path.
+///
+/// Input: `Sort(DESC) ← CoalescePartitionsExec ← multi-partition`. The first
+/// pass must rewrite this into a parallel plan
+/// `SortPreservingMergeExec ← Sort(preserve=true) ← multi-partition`, and
+/// subsequent passes must keep that parallel shape. Under the old two-rule
+/// pipeline `pushdown_sorts` could regress this back into a serial sort.
+///
+/// Runs 3 passes (one more than the standard idempotency check) and asserts
+/// the parallel plan structure survives each one.
+#[test]
+fn test_issue_21973_parallel_sort_survives_multiple_passes() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(8));
+ let coalesce: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(source));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(sort_expr, coalesce));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ let config = ConfigOptions::default();
+
+ let p1 = EnsureRequirements::new()
+ .optimize(Arc::clone(&limit), &config)
+ .expect("pass 1");
+ let s1 = plan_string(&p1);
+
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("pass 2");
+ let s2 = plan_string(&p2);
+
+ let p3 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p2), &config)
+ .expect("pass 3");
+ let s3 = plan_string(&p3);
+
+ // The parallel-sort shape must appear after pass 1 and survive every
+ // subsequent pass. Specifically: SortPreservingMergeExec on top of a
+ // Sort with preserve_partitioning=true (no CoalescePartitionsExec).
+ for (i, plan_str) in [&s1, &s2, &s3].iter().enumerate() {
+ assert!(
+ plan_str.contains("SortPreservingMergeExec"),
+ "pass {} regressed to serial: missing SortPreservingMergeExec:\n{plan_str}",
+ i + 1
+ );
+ assert!(
+ plan_str.contains("preserve_partitioning=[true]"),
+ "pass {} regressed to serial: Sort lost preserve_partitioning=true (#21973):\n{plan_str}",
+ i + 1
+ );
+ assert!(
+ !plan_str.contains("CoalescePartitionsExec"),
+ "pass {} regressed to serial: CoalescePartitionsExec re-introduced (#21973):\n{plan_str}",
+ i + 1
+ );
+ }
+
+ assert_eq!(
+ s1, s2,
+ "not idempotent between pass 1 and 2 (#21973):\n{s1}\nvs\n{s2}"
+ );
+ assert_eq!(
+ s2, s3,
+ "not idempotent between pass 2 and 3 (#21973):\n{s2}\nvs\n{s3}"
+ );
+
+ // All passes must produce sanity-checkable plans.
+ for (i, p) in [p1, p2, p3].into_iter().enumerate() {
+ SanityCheckPlan::new()
+ .optimize(p, &config)
+ .unwrap_or_else(|e| {
+ panic!("SanityCheckPlan failed on pass {}: {e:?}", i + 1)
+ });
+ }
+}
+
+/// Idempotency: Sort → Aggregate → Sort → Aggregate pattern (#18989).
+/// This tests the multi-aggregate topology that caused SanityCheckPlan
+/// failures in upstream issue #18989.
+#[test]
+fn test_idempotent_sort_aggregate_sort_aggregate() {
+ use datafusion_physical_plan::aggregates::{
+ AggregateExec, AggregateMode, PhysicalGroupBy,
+ };
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+
+ // Partial aggregate
+ let group_by = PhysicalGroupBy::new_single(vec![(
+ Arc::new(Column::new("a", 0)) as _,
+ "a".to_string(),
+ )]);
+ let partial_agg: Arc<dyn ExecutionPlan> = Arc::new(
+ AggregateExec::try_new(
+ AggregateMode::Partial,
+ group_by,
+ vec![],
+ vec![],
+ source,
+ Arc::clone(&schema),
+ )
+ .unwrap(),
+ );
+
+ let sort_expr = sort_expr_on("a", 0, false, false);
+
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(sort_expr, partial_agg));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_idempotent(limit);
+}
+
+/// Stress test: idempotency with ALL partition counts from 1 to 64
+#[test]
+fn test_idempotent_all_partition_counts_1_to_64() {
+ for n in 1..=64 {
+ let source = Arc::new(MockMultiPartitionExec::new(n));
+ let sort_expr = sort_expr_on("a", 0, true, true);
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+ assert_idempotent(limit);
+ }
+}
+
+/// Regression for #14150: the standalone distribution enforcement path
+/// lost `fetch` when applied twice. Verify `EnsureRequirements`
+/// preserves fetch across multiple passes.
+#[test]
+fn test_issue_14150_fetch_survives_multiple_passes() {
+ // Simulate: SELECT * FROM multi_partition_table ORDER BY a LIMIT 5
+ // with target_partitions > 1 (triggers RepartitionExec)
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(1));
+ let repartition = Arc::new(
+ RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(4)).unwrap(),
+ );
+
+ let sort_expr = sort_expr_on("a", 0, false, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, repartition as _).with_fetch(Some(5)));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ let config = ConfigOptions::default();
+
+ // Pass 1
+ let p1 = EnsureRequirements::new()
+ .optimize(Arc::clone(&limit), &config)
+ .unwrap();
+ let s1 = plan_string(&p1);
+
+ // Pass 2
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .unwrap();
+ let s2 = plan_string(&p2);
+
+ // Pass 3
+ let p3 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p2), &config)
+ .unwrap();
+ let s3 = plan_string(&p3);
+
+ // Fetch must survive all passes
+ assert!(s1.contains("fetch=5"), "fetch=5 lost after pass 1:\n{s1}");
+ assert!(
+ s2.contains("fetch=5"),
+ "fetch=5 lost after pass 2 (#14150 regression):\n{s2}"
+ );
+ assert!(s3.contains("fetch=5"), "fetch=5 lost after pass 3:\n{s3}");
+
+ // Plans must be identical (idempotent)
+ assert_eq!(s1, s2, "Plan changed between pass 1 and 2:\n{s1}\nvs\n{s2}");
+ assert_eq!(s2, s3, "Plan changed between pass 2 and 3:\n{s2}\nvs\n{s3}");
+}
+
+/// Sharper #14150 reproduce: input plan already contains a
+/// `SortPreservingMergeExec` with an explicit `fetch`, sitting directly
+/// above a `SortExec(fetch=…)` on a multi-partition source. This is the
+/// exact shape that originally triggered the bug — the old
+/// `EnforceDistribution::optimize` path would call
+/// `remove_dist_changing_operators()` on this SPM, strip it, and then
+/// `add_merge_on_top()` re-create an SPM **without** copying the saved
+/// `fetch`. Pass 2 saw an SPM with no fetch and #14150 silently bit.
+///
+/// `EnsureRequirements` preserves the `fetch` value across every pass.
+/// Note: it may legitimately deduplicate the `fetch` field between
+/// adjacent operators (e.g. push it onto the surrounding
+/// `GlobalLimitExec` and drop it from the SPM), so this test asserts
+/// the #14150 property — \"`fetch=5` must appear somewhere in the
+/// plan after every pass\" — rather than byte-identical idempotency
+/// (which is covered by `test_issue_14150_fetch_survives_multiple_passes`
+/// on the more realistic input shape where the SPM is inserted by the
+/// optimizer itself).
+#[test]
+fn test_issue_14150_fetch_survives_with_input_spm() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+
+ let sort_expr = sort_expr_on("a", 0, false, true);
+
+ // Sort with fetch=5 (TopK).
+ let sort = Arc::new(
+ SortExec::new(sort_expr.clone(), Arc::clone(&source)).with_fetch(Some(5)),
+ );
+
+ // SPM with fetch=5 above the sort — this is what `EnforceDistribution`
+ // used to strip and re-add without `fetch`.
+ let spm: Arc<dyn ExecutionPlan> =
+ Arc::new(SortPreservingMergeExec::new(sort_expr, sort).with_fetch(Some(5)));
+
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(spm, 0, Some(5)));
+
+ let config = ConfigOptions::default();
+
+ let p1 = EnsureRequirements::new()
+ .optimize(Arc::clone(&limit), &config)
+ .unwrap();
+ let s1 = plan_string(&p1);
+
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .unwrap();
+ let s2 = plan_string(&p2);
+
+ // The #14150 property: `fetch=5` must survive both passes (the
+ // historical bug was that pass 2 dropped it when the SPM got
+ // re-created in `add_merge_on_top`).
+ assert!(s1.contains("fetch=5"), "fetch=5 lost after pass 1:\n{s1}");
+ assert!(
+ s2.contains("fetch=5"),
+ "fetch=5 lost after pass 2 (#14150 regression):\n{s2}"
+ );
+}
+
+// ========================================================================
+// Mock operator with configurable distribution / ordering requirements
+// (used by window-function and distribution tests below)
+// ========================================================================
+
+/// Mock operator requiring specific distribution and/or ordering from its
+/// single child. Simulates operators like `BoundedWindowAggExec` that
+/// demand hash-partitioning + ordering without pulling in complex window
+/// expression machinery.
+#[derive(Debug)]
+struct MockReqExec {
+ input: Arc<dyn ExecutionPlan>,
+ dist: Distribution,
+ ord: Option<LexOrdering>,
+ properties: Arc<PlanProperties>,
+}
+
+impl MockReqExec {
+ fn new(
+ input: Arc<dyn ExecutionPlan>,
+ dist: Distribution,
+ ord: Option<LexOrdering>,
+ ) -> Self {
+ let properties = Arc::new(PlanProperties::new(
+ input.equivalence_properties().clone(),
+ input.output_partitioning().clone(),
+ input.pipeline_behavior(),
+ input.boundedness(),
+ ));
+ Self {
+ input,
+ dist,
+ ord,
+ properties,
+ }
+ }
+}
+
+impl DisplayAs for MockReqExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockReqExec")
+ }
+}
+
+impl ExecutionPlan for MockReqExec {
+ fn name(&self) -> &str {
+ "MockReqExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+ fn required_input_distribution(&self) -> Vec<Distribution> {
+ vec![self.dist.clone()]
+ }
+ fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+ vec![
+ self.ord
+ .as_ref()
+ .map(|o| OrderingRequirements::from(o.clone())),
+ ]
+ }
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![true]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ mut c: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ assert_eq!(c.len(), 1);
+ Ok(Arc::new(MockReqExec::new(
+ c.pop().expect("1 child"),
+ self.dist.clone(),
+ self.ord.clone(),
+ )))
+ }
+ fn execute(
+ &self,
+ _p: usize,
+ _c: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+}
+
+// ========================================================================
+// Additional idempotency tests covering remaining sub-passes
+// ========================================================================
+
+/// Idempotency: plan that triggers `parallelize_sorts`.
+/// CoalescePartitionsExec → SortExec(preserve=false) → multi-partition
+/// source. After optimization Sort+SPM should be parallel. Running twice
+/// must produce the same plan.
+#[test]
+fn test_idempotent_parallelize_sorts() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(8));
+ let coalesce: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(source));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ // Sort without preserve_partitioning on top of coalesced input
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(sort_expr, coalesce));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ // First pass should parallelize the sort (Sort+SPM replaces Coalesce+Sort)
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(Arc::clone(&limit), &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = plan_string(&p1);
+ let s2 = plan_string(&p2);
+ assert_eq!(
+ s1, s2,
+ "parallelize_sorts NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+}
+
+/// Idempotency: SortMergeJoinExec with two multi-partition inputs + ORDER BY
+/// + LIMIT. Tests that join key reordering + sort enforcement is stable.
+#[test]
+fn test_idempotent_sort_merge_join() {
+ let left: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+ let right: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+
+ let on = vec![(
+ Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
+ Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
+ )];
+
+ let join: Arc<dyn ExecutionPlan> = Arc::new(
+ SortMergeJoinExec::try_new(
+ left,
+ right,
+ on,
+ None,
+ JoinType::Inner,
+ vec![SortOptions {
+ descending: false,
+ nulls_first: false,
+ }],
+ NullEquality::NullEqualsNothing,
+ )
+ .expect("SortMergeJoinExec creation failed"),
+ );
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, join));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(50)));
+
+ assert_idempotent(limit);
+}
+
+/// Idempotency: window-function-like operator over multi-partition source.
+/// Uses MockReqExec with hash distribution + ordering to simulate
+/// BoundedWindowAggExec requirements. Tests that window partitioning +
+/// sort requirements are stable across optimizer passes.
+#[test]
+fn test_idempotent_window_over_multi_partition() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(8));
+
+ // Window function requires Hash(a) distribution + ordering [a ASC, b ASC]
+ let ord = LexOrdering::new(vec![
+ PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ ),
+ PhysicalSortExpr::new(
+ Arc::new(Column::new("b", 1)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ ),
+ ])
+ .unwrap();
+
+ let dist = Distribution::HashPartitioned(vec![Arc::new(Column::new("a", 0))]);
+ let window_like: Arc<dyn ExecutionPlan> =
+ Arc::new(MockReqExec::new(source, dist, Some(ord)));
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, window_like));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(20)));
+
+ assert_idempotent(limit);
+}
+
+/// Idempotency: multiple levels of sort + limit.
+/// GlobalLimitExec → SortExec → ProjectionExec → GlobalLimitExec → SortExec
+/// → multi-partition source. Tests deeply nested sort/limit stability.
+#[test]
+fn test_idempotent_nested_subqueries_sort_limit() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(8));
+
+ // Inner sort (a DESC) + inner limit — use DESC to avoid matching
+ // MockMultiPartitionExec's built-in ASC ordering, which would cause
+ // the optimizer to eliminate the sort differently across passes.
+ let inner_sort: Arc<dyn ExecutionPlan> =
+ Arc::new(SortExec::new(sort_expr_on("a", 0, true, true), source));
+ let inner_limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(inner_sort, 0, Some(100)));
+
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection: Arc<dyn ExecutionPlan> =
+ Arc::new(ProjectionExec::try_new(proj_exprs, inner_limit).unwrap());
+
+ // Outer sort (a DESC) + outer limit
+ let outer_sort: Arc<dyn ExecutionPlan> =
+ Arc::new(SortExec::new(sort_expr_on("a", 0, true, true), projection));
+ let outer_limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(outer_sort, 0, Some(10)));
+
+ assert_idempotent(outer_limit);
+}
+
+/// Idempotency: RepartitionExec(Hash) + sort + limit.
+/// Tests that hash distribution + ordering enforcement is stable.
+#[test]
+fn test_idempotent_repartition_hash_sort_limit() {
+ let source: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(8));
+
+ let hash_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![Arc::new(Column::new("a", 0))];
+ let repartition = Arc::new(
+ RepartitionExec::try_new(source, Partitioning::Hash(hash_exprs, 4)).unwrap(),
+ );
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(sort_expr, repartition));
+ let limit: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(15)));
+
+ assert_idempotent(limit);
+}
+
+/// EnsureRequirements applied twice on a HashJoinExec plan must produce
+/// identical plans. Tests that hash distribution enforcement is stable.
+#[test]
+fn test_enforce_distribution_idempotent_hash_join() {
+ let left: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+ let right: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(4));
+
+ let on = vec![(
+ Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
+ Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>,
+ )];
+
+ let join: Arc<dyn ExecutionPlan> = Arc::new(
+ HashJoinExec::try_new(
+ left,
+ right,
+ on,
+ None,
+ &JoinType::Inner,
+ None,
+ PartitionMode::Partitioned,
+ NullEquality::NullEqualsNothing,
+ false,
+ )
+ .expect("HashJoinExec creation failed"),
+ );
+
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(Arc::clone(&join), &config)
+ .expect("first EnsureRequirements pass failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second EnsureRequirements pass failed");
+
+ let s1 = plan_string(&p1);
+ let s2 = plan_string(&p2);
+
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements not idempotent for HashJoinExec!\nPass 1:\n{s1}\nPass 2:\n{s2}"
+ );
+}
+
+/// Idempotency on a complex plan:
+/// `GlobalLimitExec → SortExec → ProjectionExec → UnionExec(multi, single)`.
+/// The union + projection + sort topology has historically been a fertile
+/// ground for non-idempotent behaviour, so we keep it as a separate
+/// idempotency test — `assert_idempotent` already proves `f(f(x)) == f(x)`,
+/// which for a deterministic optimiser is equivalent to stability across
+/// any finite number of passes (the previous 10x sweep was overkill).
+#[test]
+fn test_idempotent_union_projection_sort() {
+ let live: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(16));
+ let hist: Arc<dyn ExecutionPlan> = Arc::new(MockMultiPartitionExec::new(1));
+ let union: Arc<dyn ExecutionPlan> = UnionExec::try_new(vec![live, hist]).unwrap();
+
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection: Arc<dyn ExecutionPlan> =
+ Arc::new(ProjectionExec::try_new(proj_exprs, union).unwrap());
+
+ let sort_expr = sort_expr_on("a", 0, true, true);
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let plan: Arc<dyn ExecutionPlan> = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ assert_idempotent(plan);
+}
diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs
index b7ba661d23..801c2f30f9 100644
--- a/datafusion/core/tests/physical_optimizer/mod.rs
+++ b/datafusion/core/tests/physical_optimizer/mod.rs
@@ -24,6 +24,7 @@ mod combine_partial_final_agg;
mod enforce_distribution;
mod enforce_sorting;
mod enforce_sorting_monotonicity;
+mod ensure_requirements;
mod filter_pushdown;
mod join_selection;
#[expect(clippy::needless_pass_by_value)]
diff --git a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
index 74e938e75e..297a92c45a 100644
--- a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
+++ b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs
@@ -35,7 +35,8 @@ use datafusion_physical_expr::{PhysicalExpr, physical_exprs_equal};
/// CombinePartialFinalAggregate optimizer rule combines the adjacent Partial and Final AggregateExecs
/// into a Single AggregateExec if their grouping exprs and aggregate exprs equal.
///
-/// This rule should be applied after the EnforceDistribution and EnforceSorting rules
+/// This rule should be applied after the `EnsureRequirements` rule (which
+/// handles both distribution and sorting enforcement).
#[derive(Default, Debug)]
pub struct CombinePartialFinalAggregate {}
diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
similarity index 85%
rename from datafusion/physical-optimizer/src/enforce_distribution.rs
rename to datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
index c522867c05..093a1ec14b 100644
--- a/datafusion/physical-optimizer/src/enforce_distribution.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
@@ -15,17 +15,23 @@
// specific language governing permissions and limitations
// under the License.
-//! EnforceDistribution optimizer rule inspects the physical plan with respect
-//! to distribution requirements and adds [`RepartitionExec`]s to satisfy them
-//! when necessary. If increasing parallelism is beneficial (and also desirable
-//! according to the configuration), this rule increases partition counts in
-//! the physical plan.
+//! Distribution enforcement helpers. The standalone `EnforceDistribution`
+//! rule that previously lived here has been retired in favour of
+//! `EnsureRequirements` (which composes distribution and sorting
+//! enforcement into a single idempotent pass). The helpers in this
+//! module — `adjust_input_keys_ordering`, `reorder_join_keys_to_inputs`,
+//! `DistributionContext`, `ensure_distribution`, … — are used directly
+//! by `EnsureRequirements`.
+//!
+//! These helpers inspect the physical plan with respect to distribution
+//! requirements and add [`RepartitionExec`]s to satisfy them when necessary.
+//! If increasing parallelism is beneficial (and also desirable according to
+//! configuration), they increase partition counts in the physical plan.
use std::any::Any;
use std::fmt::Debug;
use std::sync::Arc;
-use crate::optimizer::PhysicalOptimizerRule;
use crate::output_requirements::OutputRequirementExec;
use crate::utils::{
add_sort_above_with_check, is_coalesce_partitions, is_repartition,
@@ -36,7 +42,7 @@ use arrow::compute::SortOptions;
use datafusion_common::config::ConfigOptions;
use datafusion_common::error::Result;
use datafusion_common::stats::Precision;
-use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::tree_node::Transformed;
use datafusion_expr::logical_plan::{Aggregate, JoinType};
use datafusion_physical_expr::expressions::{Column, NoOp};
use datafusion_physical_expr::utils::map_columns_before_projection;
@@ -63,174 +69,11 @@ use datafusion_physical_plan::{Distribution, ExecutionPlan, Partitioning};
use itertools::izip;
-/// The `EnforceDistribution` rule ensures that distribution requirements are
-/// met. In doing so, this rule will increase the parallelism in the plan by
-/// introducing repartitioning operators to the physical plan.
-///
-/// For example, given an input such as:
-///
-///
-/// ```text
-/// ┌─────────────────────────────────┐
-/// │ │
-/// │ ExecutionPlan │
-/// │ │
-/// └─────────────────────────────────┘
-/// ▲ ▲
-/// │ │
-/// ┌─────┘ └─────┐
-/// │ │
-/// │ │
-/// │ │
-/// ┌───────────┐ ┌───────────┐
-/// │ │ │ │
-/// │ batch A1 │ │ batch B1 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A2 │ │ batch B2 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A3 │ │ batch B3 │
-/// │ │ │ │
-/// └───────────┘ └───────────┘
-///
-/// Input Input
-/// A B
-/// ```
-///
-/// This rule will attempt to add a `RepartitionExec` to increase parallelism
-/// (to 3, in this case) and create the following arrangement:
-///
-/// ```text
-/// ┌─────────────────────────────────┐
-/// │ │
-/// │ ExecutionPlan │
-/// │ │
-/// └─────────────────────────────────┘
-/// ▲ ▲ ▲ Input now has 3
-/// │ │ │ partitions
-/// ┌───────┘ │ └───────┐
-/// │ │ │
-/// │ │ │
-/// ┌───────────┐ ┌───────────┐ ┌───────────┐
-/// │ │ │ │ │ │
-/// │ batch A1 │ │ batch A3 │ │ batch B3 │
-/// │ │ │ │ │ │
-/// ├───────────┤ ├───────────┤ ├───────────┤
-/// │ │ │ │ │ │
-/// │ batch B2 │ │ batch B1 │ │ batch A2 │
-/// │ │ │ │ │ │
-/// └───────────┘ └───────────┘ └───────────┘
-/// ▲ ▲ ▲
-/// │ │ │
-/// └─────────┐ │ ┌──────────┘
-/// │ │ │
-/// │ │ │
-/// ┌─────────────────────────────────┐ batches are
-/// │ RepartitionExec(3) │ repartitioned
-/// │ RoundRobin │
-/// │ │
-/// └─────────────────────────────────┘
-/// ▲ ▲
-/// │ │
-/// ┌─────┘ └─────┐
-/// │ │
-/// │ │
-/// │ │
-/// ┌───────────┐ ┌───────────┐
-/// │ │ │ │
-/// │ batch A1 │ │ batch B1 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A2 │ │ batch B2 │
-/// │ │ │ │
-/// ├───────────┤ ├───────────┤
-/// │ │ │ │
-/// │ batch A3 │ │ batch B3 │
-/// │ │ │ │
-/// └───────────┘ └───────────┘
-///
-///
-/// Input Input
-/// A B
-/// ```
-///
-/// The `EnforceDistribution` rule
-/// - is idempotent; i.e. it can be applied multiple times, each time producing
-/// the same result.
-/// - always produces a valid plan in terms of distribution requirements. Its
-/// input plan can be valid or invalid with respect to distribution requirements,
-/// but the output plan will always be valid.
-/// - produces a valid plan in terms of ordering requirements, *if* its input is
-/// a valid plan in terms of ordering requirements. If the input plan is invalid,
-/// this rule does not attempt to fix it as doing so is the responsibility of the
-/// `EnforceSorting` rule.
-///
-/// Note that distribution requirements are met in the strictest way. This may
-/// result in more than strictly necessary [`RepartitionExec`]s in the plan, but
-/// meeting the requirements in the strictest way may help avoid possible data
-/// skew in joins.
-///
-/// For example for a hash join with keys (a, b, c), the required Distribution(a, b, c)
-/// can be satisfied by several alternative partitioning ways: (a, b, c), (a, b),
-/// (a, c), (b, c), (a), (b), (c) and ( ).
-///
-/// This rule only chooses the exact match and satisfies the Distribution(a, b, c)
-/// by a HashPartition(a, b, c).
-#[derive(Default, Debug)]
-pub struct EnforceDistribution {}
-
-impl EnforceDistribution {
- #[expect(missing_docs)]
- pub fn new() -> Self {
- Self {}
- }
-}
-
-impl PhysicalOptimizerRule for EnforceDistribution {
- fn optimize(
- &self,
- plan: Arc<dyn ExecutionPlan>,
- config: &ConfigOptions,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
-
- let adjusted = if top_down_join_key_reordering {
- // Run a top-down process to adjust input key ordering recursively
- let plan_requirements = PlanWithKeyRequirements::new_default(plan);
- let adjusted = plan_requirements
- .transform_down(adjust_input_keys_ordering)
- .data()?;
- adjusted.plan
- } else {
- // Run a bottom-up process
- plan.transform_up(|plan| {
- Ok(Transformed::yes(reorder_join_keys_to_inputs(plan)?))
- })
- .data()?
- };
-
- let distribution_context = DistributionContext::new_default(adjusted);
- // Distribution enforcement needs to be applied bottom-up.
- let distribution_context = distribution_context
- .transform_up(|distribution_context| {
- ensure_distribution(distribution_context, config)
- })
- .data()?;
- Ok(distribution_context.plan)
- }
-
- fn name(&self) -> &str {
- "EnforceDistribution"
- }
-
- fn schema_check(&self) -> bool {
- true
- }
-}
+// The `EnforceDistribution` rule was retired in favour of `EnsureRequirements`,
+// which composes distribution and sorting enforcement into a single idempotent
+// pass. The helper functions below (`adjust_input_keys_ordering`,
+// `reorder_join_keys_to_inputs`, `DistributionContext`, `ensure_distribution`,
+// etc.) remain — `EnsureRequirements` calls into them directly.
#[derive(Debug, Clone)]
struct JoinKeyPairs {
@@ -970,7 +813,10 @@ fn preserving_order_enables_streaming(
///
/// Updated node with an execution plan, where the desired single distribution
/// requirement is satisfied.
-fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
+fn add_merge_on_top(
+ input: DistributionContext,
+ fetch: Option<usize>,
+) -> DistributionContext {
// Apply only when the partition count is larger than one.
if input.plan.output_partitioning().partition_count() > 1 {
// When there is an existing ordering, we preserve ordering
@@ -979,14 +825,20 @@ fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
// - Preserving ordering is not helpful in terms of satisfying ordering requirements
// - Usage of order preserving variants is not desirable
// (determined by flag `config.optimizer.prefer_existing_sort`)
- let new_plan = if let Some(req) = input.plan.output_ordering() {
- Arc::new(SortPreservingMergeExec::new(
- req.clone(),
- Arc::clone(&input.plan),
- )) as _
+ let new_plan: Arc<dyn ExecutionPlan> = if let Some(req) =
+ input.plan.output_ordering()
+ {
+ let mut spm =
+ SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan));
+ if let Some(f) = fetch {
+ spm = spm.with_fetch(Some(f));
+ }
+ Arc::new(spm)
} else {
// If there is no input order, we can simply coalesce partitions:
- Arc::new(CoalescePartitionsExec::new(Arc::clone(&input.plan))) as _
+ Arc::new(
+ CoalescePartitionsExec::new(Arc::clone(&input.plan)).with_fetch(fetch),
+ )
};
DistributionContext::new(new_plan, true, vec![input])
@@ -1012,20 +864,41 @@ fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
/// ```text
/// "DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet",
/// ```
+/// Returned by [`remove_dist_changing_operators`] to carry the fetch value
+/// that may have been on a removed `SortPreservingMergeExec` or `CoalescePartitionsExec`.
+struct RemovedDistOps {
+ context: DistributionContext,
+ /// The fetch value from the removed SPM/Coalesce, if any.
+ /// Must be re-applied when distribution operators are re-inserted.
+ removed_fetch: Option<usize>,
+}
+
fn remove_dist_changing_operators(
mut distribution_context: DistributionContext,
-) -> Result<DistributionContext> {
+) -> Result<RemovedDistOps> {
+ let mut removed_fetch = None;
while is_repartition(&distribution_context.plan)
|| is_coalesce_partitions(&distribution_context.plan)
|| is_sort_preserving_merge(&distribution_context.plan)
{
+ // Preserve fetch from SPM or CoalescePartitions before removing (#14150).
+ if let Some(fetch) = distribution_context.plan.fetch() {
+ removed_fetch = Some(
+ removed_fetch
+ .map(|existing: usize| existing.min(fetch))
+ .unwrap_or(fetch),
+ );
+ }
// All of above operators have a single child. First child is only child.
// Remove any distribution changing operators at the beginning:
distribution_context = distribution_context.children.swap_remove(0);
// Note that they will be re-inserted later on if necessary or helpful.
}
- Ok(distribution_context)
+ Ok(RemovedDistOps {
+ context: distribution_context,
+ removed_fetch,
+ })
}
/// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable.
@@ -1219,11 +1092,16 @@ pub fn ensure_distribution(
let order_preserving_variants_desirable =
unbounded_and_pipeline_friendly || config.optimizer.prefer_existing_sort;
- // Remove unnecessary repartition from the physical plan if any
- let DistributionContext {
- mut plan,
- data,
- children,
+ // Remove unnecessary repartition from the physical plan if any.
+ // Preserve fetch from removed SPM/Coalesce (#14150).
+ let RemovedDistOps {
+ context:
+ DistributionContext {
+ mut plan,
+ data,
+ children,
+ },
+ removed_fetch,
} = remove_dist_changing_operators(dist_context)?;
if let Some(exec) = plan.downcast_ref::<WindowAggExec>() {
@@ -1359,7 +1237,7 @@ pub fn ensure_distribution(
// Satisfy the distribution requirement if it is unmet.
match &requirement {
Distribution::SinglePartition => {
- child = add_merge_on_top(child);
+ child = add_merge_on_top(child, removed_fetch);
}
Distribution::HashPartitioned(exprs) => {
// See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background
diff --git a/datafusion/physical-optimizer/src/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs
similarity index 87%
rename from datafusion/physical-optimizer/src/enforce_sorting/mod.rs
rename to datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs
index 241c843556..53917a5108 100644
--- a/datafusion/physical-optimizer/src/enforce_sorting/mod.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs
@@ -15,38 +15,39 @@
// specific language governing permissions and limitations
// under the License.
-//! EnforceSorting optimizer rule inspects the physical plan with respect
-//! to local sorting requirements and does the following:
+//! Sort enforcement helpers. The standalone `EnforceSorting` rule that
+//! previously lived here has been retired in favour of `EnsureRequirements`
+//! (which composes distribution and sorting enforcement into a single
+//! idempotent pass). The helpers in this module — `ensure_sorting`,
+//! `parallelize_sorts`, `PlanWithCorrespondingSort`, and the submodules
+//! `replace_with_order_preserving_variants` and `sort_pushdown` — are
+//! used directly by `EnsureRequirements`.
+//!
+//! Sort enforcement inspects the physical plan with respect to local
+//! sorting requirements and does the following:
//! - Adds a [`SortExec`] when a requirement is not met,
//! - Removes an already-existing [`SortExec`] if it is possible to prove
//! that this sort is unnecessary
//!
-//! The rule can work on valid *and* invalid physical plans with respect to
-//! sorting requirements, but always produces a valid physical plan in this sense.
+//! The helpers can work on valid *and* invalid physical plans with respect
+//! to sorting requirements, but always produce a valid plan in this sense.
//!
-//! A non-realistic but easy to follow example for sort removals: Assume that we
-//! somehow get the fragment
+//! A non-realistic but easy to follow example for sort removals: assume the
+//! fragment
//!
//! ```text
//! SortExec: expr=[nullable_col@0 ASC]
//! SortExec: expr=[non_nullable_col@1 ASC]
//! ```
//!
-//! in the physical plan. The first sort is unnecessary since its result is overwritten
-//! by another [`SortExec`]. Therefore, this rule removes it from the physical plan.
+//! reaches this stage. The first sort is unnecessary since its result is
+//! overwritten by another [`SortExec`], so it is removed.
pub mod replace_with_order_preserving_variants;
pub mod sort_pushdown;
use std::sync::Arc;
-use crate::PhysicalOptimizerRule;
-use crate::enforce_sorting::replace_with_order_preserving_variants::{
- OrderPreservationContext, replace_with_order_preserving_variants,
-};
-use crate::enforce_sorting::sort_pushdown::{
- SortPushDown, assign_initial_requirements, pushdown_sorts,
-};
use crate::output_requirements::OutputRequirementExec;
use crate::utils::{
add_sort_above, add_sort_above_with_check, is_coalesce_partitions, is_limit,
@@ -54,9 +55,8 @@ use crate::utils::{
};
use datafusion_common::Result;
-use datafusion_common::config::ConfigOptions;
use datafusion_common::plan_err;
-use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::tree_node::Transformed;
use datafusion_physical_expr::{Distribution, Partitioning};
use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement};
use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
@@ -73,19 +73,13 @@ use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties, InputOrde
use itertools::izip;
-/// This rule inspects [`SortExec`]'s in the given physical plan in order to
-/// remove unnecessary sorts, and optimize sort performance across the plan.
-#[derive(Default, Debug)]
-pub struct EnforceSorting {}
-
-impl EnforceSorting {
- #[expect(missing_docs)]
- pub fn new() -> Self {
- Self {}
- }
-}
+// The `EnforceSorting` rule was retired in favour of `EnsureRequirements`,
+// which composes distribution and sorting enforcement into a single idempotent
+// pass. The helper functions and contexts below (`ensure_sorting`,
+// `parallelize_sorts`, `PlanWithCorrespondingSort`, etc.) remain —
+// `EnsureRequirements` calls into them directly.
-/// This context object is used within the [`EnforceSorting`] rule to track the closest
+/// Context object used by sort enforcement to track the closest
/// [`SortExec`] descendant(s) for every child of a plan. The data attribute
/// stores whether the plan is a `SortExec` or is connected to a `SortExec`
/// via its children.
@@ -135,7 +129,7 @@ fn update_sort_ctx_children_data(
Ok(node_and_ctx)
}
-/// This object is used within the [`EnforceSorting`] rule to track the closest
+/// Tracks the closest
/// [`CoalescePartitionsExec`] descendant(s) for every child of a plan. The data
/// attribute stores whether the plan is a `CoalescePartitionsExec` or is
/// connected to a `CoalescePartitionsExec` via its children.
@@ -191,78 +185,11 @@ fn update_coalesce_ctx_children(
};
}
-/// Performs optimizations based upon a series of subrules.
-/// Refer to each subrule for detailed descriptions of the optimizations performed:
-/// Subrule application is ordering dependent.
-///
-/// Optimizer consists of 5 main parts which work sequentially
-/// 1. [`ensure_sorting`] Works down-to-top to be able to remove unnecessary [`SortExec`]s, [`SortPreservingMergeExec`]s
-/// add [`SortExec`]s if necessary by a requirement and adjusts window operators.
-/// 2. [`parallelize_sorts`] (Optional, depends on the `repartition_sorts` configuration)
-/// Responsible to identify and remove unnecessary partition unifier operators
-/// such as [`SortPreservingMergeExec`], [`CoalescePartitionsExec`] follows [`SortExec`]s does possible simplifications.
-/// 3. [`replace_with_order_preserving_variants()`] Replaces with alternative operators, for example can merge
-/// a [`SortExec`] and a [`CoalescePartitionsExec`] into one [`SortPreservingMergeExec`]
-/// or a [`SortExec`] + [`RepartitionExec`] combination into an order preserving [`RepartitionExec`]
-/// 4. [`sort_pushdown`] Works top-down. Responsible to push down sort operators as deep as possible in the plan.
-/// 5. `replace_with_partial_sort` Checks if it's possible to replace [`SortExec`]s with [`PartialSortExec`] operators
-impl PhysicalOptimizerRule for EnforceSorting {
- fn optimize(
- &self,
- plan: Arc<dyn ExecutionPlan>,
- config: &ConfigOptions,
- ) -> Result<Arc<dyn ExecutionPlan>> {
- let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
- // Execute a bottom-up traversal to enforce sorting requirements,
- // remove unnecessary sorts, and optimize sort-sensitive operators:
- let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
- let new_plan = if config.optimizer.repartition_sorts {
- let plan_with_coalesce_partitions =
- PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
- let parallel = plan_with_coalesce_partitions
- .transform_up(parallelize_sorts)
- .data()?;
- parallel.plan
- } else {
- adjusted.plan
- };
-
- let plan_with_pipeline_fixer = OrderPreservationContext::new_default(new_plan);
- let updated_plan = plan_with_pipeline_fixer
- .transform_up(|plan_with_pipeline_fixer| {
- replace_with_order_preserving_variants(
- plan_with_pipeline_fixer,
- false,
- true,
- config,
- )
- })
- .data()?;
- // Execute a top-down traversal to exploit sort push-down opportunities
- // missed by the bottom-up traversal:
- let mut sort_pushdown = SortPushDown::new_default(updated_plan.plan);
- assign_initial_requirements(&mut sort_pushdown);
- let adjusted = pushdown_sorts(sort_pushdown)?;
- adjusted
- .plan
- .transform_up(|plan| Ok(Transformed::yes(replace_with_partial_sort(plan)?)))
- .data()
- }
-
- fn name(&self) -> &str {
- "EnforceSorting"
- }
-
- fn schema_check(&self) -> bool {
- true
- }
-}
-
/// Only interested with [`SortExec`]s and their unbounded children.
/// If the plan is not a [`SortExec`] or its child is not unbounded, returns the original plan.
/// Otherwise, by checking the requirement satisfaction searches for a replacement chance.
/// If there's one replaces the [`SortExec`] plan with a [`PartialSortExec`]
-fn replace_with_partial_sort(
+pub fn replace_with_partial_sort(
plan: Arc<dyn ExecutionPlan>,
) -> Result<Arc<dyn ExecutionPlan>> {
let Some(sort_plan) = plan.downcast_ref::<SortExec>() else {
diff --git a/datafusion/physical-optimizer/src/enforce_sorting/replace_with_order_preserving_variants.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs
similarity index 100%
rename from datafusion/physical-optimizer/src/enforce_sorting/replace_with_order_preserving_variants.rs
rename to datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs
diff --git a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs
similarity index 87%
rename from datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs
rename to datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs
index 400161a94c..261cf701c8 100644
--- a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs
@@ -19,7 +19,8 @@ use std::fmt::Debug;
use std::sync::Arc;
use crate::utils::{
- add_sort_above, is_sort, is_sort_preserving_merge, is_union, is_window,
+ add_sort_above_with_distribution, is_sort, is_sort_preserving_merge, is_union,
+ is_window,
};
use arrow::datatypes::SchemaRef;
@@ -29,7 +30,7 @@ use datafusion_expr::JoinType;
use datafusion_physical_expr::expressions::Column;
use datafusion_physical_expr::utils::collect_columns;
use datafusion_physical_expr::{
- EquivalenceProperties, add_offset_to_physical_sort_exprs,
+ Distribution, EquivalenceProperties, add_offset_to_physical_sort_exprs,
};
use datafusion_physical_expr_common::sort_expr::{
LexOrdering, LexRequirement, OrderingRequirements, PhysicalSortExpr,
@@ -48,17 +49,30 @@ use datafusion_physical_plan::sorts::sort::SortExec;
use datafusion_physical_plan::tree_node::PlanContext;
use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
-/// This is a "data class" we use within the [`EnforceSorting`] rule to push
-/// down [`SortExec`] in the plan. In some cases, we can reduce the total
-/// computational cost by pushing down `SortExec`s through some executors. The
-/// object carries the parent required ordering and the (optional) `fetch` value
-/// of the parent node as its data.
-///
-/// [`EnforceSorting`]: crate::enforce_sorting::EnforceSorting
-#[derive(Default, Clone, Debug)]
+/// "Data class" used by sort pushdown (now driven from `EnsureRequirements`)
+/// to push down [`SortExec`] in the plan. In some cases the total
+/// computational cost is reduced by pushing down `SortExec`s through certain
+/// executors. The object carries the parent required ordering, the (optional)
+/// `fetch` value of the parent node, and the parent's distribution requirement
+/// (used by the distribution-aware pushdown path) as its data.
+#[derive(Clone, Debug)]
pub struct ParentRequirements {
ordering_requirement: Option<OrderingRequirements>,
fetch: Option<usize>,
+ /// The distribution required by the consumer above any SortExec we insert.
+ /// When this is `SinglePartition` and the input has multiple partitions,
+ /// `add_sort_above_with_distribution` wraps the sort in `SortPreservingMergeExec`.
+ distribution_requirement: Distribution,
+}
+
+impl Default for ParentRequirements {
+ fn default() -> Self {
+ Self {
+ ordering_requirement: None,
+ fetch: None,
+ distribution_requirement: Distribution::UnspecifiedDistribution,
+ }
+ }
}
pub type SortPushDown = PlanContext<ParentRequirements>;
@@ -66,12 +80,17 @@ pub type SortPushDown = PlanContext<ParentRequirements>;
/// Assigns the ordering requirement of the root node to the its children.
pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) {
let reqs = sort_push_down.plan.required_input_ordering();
- for (child, requirement) in sort_push_down.children.iter_mut().zip(reqs) {
+ let dists = sort_push_down.plan.required_input_distribution();
+ for (idx, (child, requirement)) in
+ sort_push_down.children.iter_mut().zip(reqs).enumerate()
+ {
child.data = ParentRequirements {
ordering_requirement: requirement,
- // If the parent has a fetch value, assign it to the children
- // Or use the fetch value of the child.
fetch: child.plan.fetch(),
+ distribution_requirement: dists
+ .get(idx)
+ .cloned()
+ .unwrap_or(Distribution::UnspecifiedDistribution),
};
}
}
@@ -92,11 +111,25 @@ fn min_fetch(f1: Option<usize>, f2: Option<usize>) -> Option<usize> {
}
}
+/// Returns the stricter of two distribution requirements.
+/// `SinglePartition` is the strictest.
+fn stronger_distribution(a: &Distribution, b: &Distribution) -> Distribution {
+ match (a, b) {
+ (Distribution::SinglePartition, _) | (_, Distribution::SinglePartition) => {
+ Distribution::SinglePartition
+ }
+ (Distribution::HashPartitioned(_), _) => a.clone(),
+ (_, Distribution::HashPartitioned(_)) => b.clone(),
+ _ => Distribution::UnspecifiedDistribution,
+ }
+}
+
fn pushdown_sorts_helper(
mut sort_push_down: SortPushDown,
) -> Result<Transformed<SortPushDown>> {
let plan = sort_push_down.plan;
let parent_fetch = sort_push_down.data.fetch;
+ let parent_distribution = sort_push_down.data.distribution_requirement.clone();
let Some(parent_requirement) = sort_push_down.data.ordering_requirement.clone()
else {
@@ -121,6 +154,14 @@ fn pushdown_sorts_helper(
return pushdown_sorts_helper(sort_push_down);
}
sort_push_down.plan = plan;
+ // No ordering is being pushed; use each child's own distribution requirement
+ let dists = sort_push_down.plan.required_input_distribution();
+ for (idx, child) in sort_push_down.children.iter_mut().enumerate() {
+ child.data.distribution_requirement = dists
+ .get(idx)
+ .cloned()
+ .unwrap_or(Distribution::UnspecifiedDistribution);
+ }
return Ok(Transformed::no(sort_push_down));
};
@@ -149,22 +190,29 @@ fn pushdown_sorts_helper(
// The sort was imposing a different ordering than the one being
// pushed down. Replace it with a sort that matches the pushed-down
// ordering, and continue the pushdown.
- // Add back the sort:
- sort_push_down = add_sort_above(
+ // Add back the sort (distribution-aware):
+ sort_push_down = add_sort_above_with_distribution(
sort_push_down,
parent_requirement.into_single(),
parent_fetch,
+ &parent_distribution,
);
// Update pushdown requirements:
sort_push_down.children[0].data = ParentRequirements {
ordering_requirement: Some(OrderingRequirements::from(sort_ordering)),
fetch: sort_fetch,
+ distribution_requirement: Distribution::UnspecifiedDistribution,
};
return Ok(Transformed::yes(sort_push_down));
} else {
// Sort was unnecessary, just propagate the stricter fetch and
- // ordering requirements:
+ // ordering requirements. Reset distribution to Unspecified
+ // because the sort we're removing may have been below a
+ // partition-merging node (like SortPreservingMergeExec) that
+ // already satisfies SinglePartition.
sort_push_down.data.fetch = min_fetch(sort_fetch, parent_fetch);
+ sort_push_down.data.distribution_requirement =
+ Distribution::UnspecifiedDistribution;
let current_is_stricter = eqp.requirements_compatible(
sort_ordering.clone().into(),
parent_requirement.first().clone(),
@@ -184,10 +232,28 @@ fn pushdown_sorts_helper(
if satisfy_parent {
// For non-sort operators which satisfy ordering:
let reqs = sort_push_down.plan.required_input_ordering();
+ let dists = sort_push_down.plan.required_input_distribution();
+
+ // If this node already outputs single partition, don't push SinglePartition
+ // requirement to children (they're below the merge point).
+ let effective_parent_dist =
+ if sort_push_down.plan.output_partitioning().partition_count() == 1 {
+ Distribution::UnspecifiedDistribution
+ } else {
+ parent_distribution.clone()
+ };
- for (child, order) in sort_push_down.children.iter_mut().zip(reqs) {
+ for (idx, (child, order)) in
+ sort_push_down.children.iter_mut().zip(reqs).enumerate()
+ {
child.data.ordering_requirement = order;
child.data.fetch = min_fetch(parent_fetch, child.data.fetch);
+ child.data.distribution_requirement = stronger_distribution(
+ &effective_parent_dist,
+ dists
+ .get(idx)
+ .unwrap_or(&Distribution::UnspecifiedDistribution),
+ );
}
} else if let Some(adjusted) = pushdown_requirement_to_children(
&sort_push_down.plan,
@@ -195,19 +261,36 @@ fn pushdown_sorts_helper(
parent_fetch,
)? {
// For operators that can take a sort pushdown, continue with updated
- // requirements:
+ // requirements. If this node already outputs single partition (e.g. SPM),
+ // don't push SinglePartition to children.
let current_fetch = sort_push_down.plan.fetch();
- for (child, order) in sort_push_down.children.iter_mut().zip(adjusted) {
+ let dists = sort_push_down.plan.required_input_distribution();
+ let effective_dist =
+ if sort_push_down.plan.output_partitioning().partition_count() == 1 {
+ Distribution::UnspecifiedDistribution
+ } else {
+ parent_distribution.clone()
+ };
+ for (idx, (child, order)) in
+ sort_push_down.children.iter_mut().zip(adjusted).enumerate()
+ {
child.data.ordering_requirement = order;
child.data.fetch = min_fetch(current_fetch, parent_fetch);
+ child.data.distribution_requirement = stronger_distribution(
+ &effective_dist,
+ dists
+ .get(idx)
+ .unwrap_or(&Distribution::UnspecifiedDistribution),
+ );
}
sort_push_down.data.ordering_requirement = None;
} else {
- // Can not push down requirements, add new `SortExec`:
- sort_push_down = add_sort_above(
+ // Can not push down requirements, add new `SortExec` (distribution-aware):
+ sort_push_down = add_sort_above_with_distribution(
sort_push_down,
parent_requirement.into_single(),
parent_fetch,
+ &parent_distribution,
);
assign_initial_requirements(&mut sort_push_down);
}
diff --git a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs
new file mode 100644
index 0000000000..41a03bb031
--- /dev/null
+++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs
@@ -0,0 +1,259 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces distribution and
+//! sorting requirements together so that the two never invalidate each other.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! `optimize` runs several tree traversals. The defining property of this
+//! rule is **Phase 2**: a single combined bottom-up pass that resolves
+//! distribution *and* sorting for each node together. The surrounding phases
+//! are independent traversals (top-down join-key reorder, then several
+//! follow-up sort/order rewrites). Some of those could be consolidated
+//! further in a follow-up.
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1: top-down join-key reorder (adjust_input_keys_ordering)
+//! │
+//! ├─ Phase 2: combined distribution + sorting (single bottom-up pass)
+//! │ └─ For each node (bottom-up), for each child:
+//! │ Step 1: ensure distribution requirement
+//! │ └─ insert RepartitionExec / CoalescePartitionsExec /
+//! │ SortPreservingMergeExec as needed
+//! │ Step 2: ensure ordering requirement (distribution-aware)
+//! │ └─ insert SortExec with the correct `preserve_partitioning`,
+//! │ with SortPreservingMergeExec on top if needed
+//! │
+//! └─ Phase 3: small follow-up passes (bottom-up unless noted)
+//! ├─ parallelize_sorts
+//! ├─ replace_with_order_preserving_variants
+//! ├─ pushdown_sorts (recursive walk)
+//! └─ replace_with_partial_sort
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent across the whole rule**: Running `EnsureRequirements`
+//! twice produces the same plan. This is the property that fixes
+//! <https://github.com/apache/datafusion/issues/21973>, where the old
+//! two-rule pipeline could regress a parallel sort plan into a serial one
+//! on pass 2.
+//! - **Distribution before sorting**: For each child, distribution is
+//! resolved before ordering, so sorting decisions always have full
+//! distribution context.
+//! - **Sort pushdown is implicit**: Phase 2 only adds `SortExec` where the
+//! child doesn't already satisfy the ordering requirement, so sorts land
+//! at the deepest valid position without a separate destructive pass.
+//!
+//! # Behavior: parallelism via repartitioning
+//!
+//! Phase 2 Step 1 inserts `RepartitionExec` to satisfy distribution
+//! requirements. When configuration allows, it also increases parallelism by
+//! repartitioning over otherwise-serial inputs. For example, given two
+//! 1-partition inputs feeding an operator that can run with more
+//! parallelism:
+//!
+//! ```text
+//! ┌─────────────────────────────────┐
+//! │ ExecutionPlan │
+//! └─────────────────────────────────┘
+//! ▲ ▲
+//! │ │
+//! ┌───────────┐ ┌───────────┐
+//! │ batch A │ │ batch B │ Input: 2 partitions
+//! └───────────┘ └───────────┘
+//! ```
+//!
+//! `EnsureRequirements` inserts a `RepartitionExec` so the operator runs
+//! with three partitions:
+//!
+//! ```text
+//! ┌─────────────────────────────────┐
+//! │ ExecutionPlan │ Input now has 3 partitions
+//! └─────────────────────────────────┘
+//! ▲ ▲ ▲
+//! └──────┼───────┘
+//! │
+//! ┌─────────────────────────────────┐
+//! │ RepartitionExec(3) │ batches are repartitioned
+//! │ RoundRobin │
+//! └─────────────────────────────────┘
+//! ▲ ▲
+//! ┌───────────┐ ┌───────────┐
+//! │ batch A │ │ batch B │
+//! └───────────┘ └───────────┘
+//! ```
+//!
+//! # Behavior: joint distribution + sorting
+//!
+//! Resolving distribution and sorting together lets Phase 2 produce a
+//! parallel sort plan in cases where the two-rule pipeline historically
+//! risked a serial one. Given `Sort(DESC) ← Coalesce ← MultiPartitionSource`,
+//! `EnsureRequirements` rewrites it into:
+//!
+//! ```text
+//! SortPreservingMergeExec: [a DESC] (cheap k-way merge of sorted streams)
+//! SortExec: [a DESC], preserve_partitioning=true (N sorts run in parallel)
+//! MultiPartitionSource
+//! ```
+//!
+//! Each input partition is sorted in parallel, then a `SortPreservingMergeExec`
+//! at the top performs a cheap merge of pre-sorted streams. For TopK queries
+//! (`fetch=K`), each parallel sort only keeps K rows per partition, so total
+//! memory is `N × K` rather than coalescing the entire stream first.
+//!
+//! # Behavior: strictest distribution match for joins
+//!
+//! Distribution requirements are met in the strictest way. For example, a
+//! hash join with keys `(a, b, c)` requires `Distribution(a, b, c)`. This
+//! can in principle be satisfied by partitioning on any superset of any
+//! subset of `(a, b, c)`, but this rule always partitions on the exact key
+//! tuple `(a, b, c)`. This is sometimes more aggressive than strictly
+//! necessary, but the strictest match helps avoid data skew in joins.
+
+// Internal implementation modules. Re-exported from `crate` root for tests
+// in `core/tests/physical_optimizer/{enforce_distribution,enforce_sorting}.rs`.
+pub mod enforce_distribution;
+pub mod enforce_sorting;
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use super::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use super::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use super::enforce_distribution::{DistributionContext, ensure_distribution};
+ use super::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use super::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use super::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use super::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+// See tests in datafusion/core/tests/physical_optimizer/ensure_requirements.rs
diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs
index 5fac8948b7..b9eb248f6e 100644
--- a/datafusion/physical-optimizer/src/lib.rs
+++ b/datafusion/physical-optimizer/src/lib.rs
@@ -27,9 +27,12 @@
pub mod aggregate_statistics;
pub mod combine_partial_final_agg;
-pub mod enforce_distribution;
-pub mod enforce_sorting;
pub mod ensure_coop;
+pub mod ensure_requirements;
+// `enforce_distribution` and `enforce_sorting` are now internal implementation
+// details of `ensure_requirements`. Re-export at the crate root so external test
+// modules keep their public paths.
+pub use ensure_requirements::{enforce_distribution, enforce_sorting};
pub mod filter_pushdown;
pub mod join_selection;
pub mod limit_pushdown;
diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs
index 05df642f84..0f81512b61 100644
--- a/datafusion/physical-optimizer/src/optimizer.rs
+++ b/datafusion/physical-optimizer/src/optimizer.rs
@@ -22,9 +22,8 @@ use std::sync::Arc;
use crate::aggregate_statistics::AggregateStatistics;
use crate::combine_partial_final_agg::CombinePartialFinalAggregate;
-use crate::enforce_distribution::EnforceDistribution;
-use crate::enforce_sorting::EnforceSorting;
use crate::ensure_coop::EnsureCooperative;
+use crate::ensure_requirements::EnsureRequirements;
use crate::filter_pushdown::FilterPushdown;
use crate::join_selection::JoinSelection;
use crate::limit_pushdown::LimitPushdown;
@@ -156,11 +155,11 @@ impl PhysicalOptimizer {
Arc::new(AggregateStatistics::new()),
// Statistics-based join selection will change the Auto mode to a real join implementation,
// like collect left, or hash join, or future sort merge join, which will influence the
- // EnforceDistribution and EnforceSorting rules as they decide whether to add additional
- // repartitioning and local sorting steps to meet distribution and ordering requirements.
- // Therefore, it should run before EnforceDistribution and EnforceSorting.
+ // EnsureRequirements rule as it decides whether to add additional repartitioning and
+ // local sorting steps to meet distribution and ordering requirements. Therefore, it
+ // should run before EnsureRequirements.
Arc::new(JoinSelection::new()),
- // The LimitedDistinctAggregation rule should be applied before the EnforceDistribution rule,
+ // The LimitedDistinctAggregation rule should be applied before EnsureRequirements,
// as that rule may inject other operations in between the different AggregateExecs.
// Applying the rule early means only directly-connected AggregateExecs must be examined.
Arc::new(LimitedDistinctAggregation::new()),
@@ -170,23 +169,32 @@ impl PhysicalOptimizer {
// those are handled by the later `FilterPushdown` rule.
// See `FilterPushdownPhase` for more details.
Arc::new(FilterPushdown::new()),
- // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
- // requirements. Please make sure that the whole plan tree is determined before this rule.
- // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at
- // least one of the operators in the plan benefits from increased parallelism.
- Arc::new(EnforceDistribution::new()),
- // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
+ // Ensures each input plan satisfies the distribution and ordering
+ // requirements declared by `ExecutionPlan::required_input_distribution`
+ // and `ExecutionPlan::required_input_ordering`.
+ //
+ // If the requirements are already satisfied, this rule leaves the plan
+ // unchanged. For example, it does not add sorting when the input is a
+ // file scan whose existing order already satisfies the required ordering.
+ // Otherwise, this rule inserts the necessary repartitioning and sorting
+ // operators.
+ //
+ // This used to be implemented as two separate rules: `EnforceDistribution`
+ // and `EnforceSorting`. It is now a single idempotent rule that decides
+ // distribution and sorting together in one bottom-up pass, so the
+ // `pushdown_sorts` step no longer breaks distribution invariants set
+ // earlier in the pipeline. See the module-level doc on
+ // [`EnsureRequirements`](crate::ensure_requirements) for the per-phase
+ // breakdown, and <https://github.com/apache/datafusion/issues/21973>
+ // for the original failure mode.
+ Arc::new(EnsureRequirements::new()),
+ // The CombinePartialFinalAggregate rule should be applied after distribution enforcement
Arc::new(CombinePartialFinalAggregate::new()),
- // The EnforceSorting rule is for adding essential local sorting to satisfy the required
- // ordering. Please make sure that the whole plan tree is determined before this rule.
- // Note that one should always run this rule after running the EnforceDistribution rule
- // as the latter may break local sorting requirements.
- Arc::new(EnforceSorting::new()),
// Run once after the local sorting requirement is changed
Arc::new(OptimizeAggregateOrder::new()),
// WindowTopN: replaces Filter(rn<=K) → Window(ROW_NUMBER) → Sort
// with Window(ROW_NUMBER) → PartitionedTopKExec(fetch=K).
- // Must run after EnforceSorting (which inserts SortExec) and before
+ // Must run after EnsureRequirements (which inserts SortExec) and before
// ProjectionPushdown (which embeds projections into FilterExec).
Arc::new(WindowTopN::new()),
// TODO: `try_embed_to_hash_join` in the ProjectionPushdown rule would be block by the CoalesceBatches, so add it before CoalesceBatches. Maybe optimize it in the future.
@@ -201,7 +209,7 @@ impl PhysicalOptimizer {
Arc::new(TopKAggregation::new()),
// Tries to push limits down through window functions, growing as appropriate
// This can possibly be combined with [LimitPushdown]
- // It needs to come after [EnforceSorting]
+ // It needs to come after [EnsureRequirements] (which handles sort enforcement)
Arc::new(LimitPushPastWindows::new()),
// The HashJoinBuffering rule adds a BufferExec node with the configured capacity
// in the prob side of hash joins. That way, the probe side gets eagerly polled before
diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs
index a6b01637c9..04229e1cc2 100644
--- a/datafusion/physical-optimizer/src/utils.rs
+++ b/datafusion/physical-optimizer/src/utils.rs
@@ -18,7 +18,7 @@
use std::sync::Arc;
use datafusion_common::Result;
-use datafusion_physical_expr::{LexOrdering, LexRequirement};
+use datafusion_physical_expr::{Distribution, LexOrdering, LexRequirement};
use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::repartition::RepartitionExec;
@@ -58,6 +58,56 @@ pub fn add_sort_above<T: Clone + Default>(
PlanContext::new(Arc::new(new_sort), T::default(), vec![node])
}
+/// Like [`add_sort_above`], but also inserts a [`SortPreservingMergeExec`] when
+/// the parent distribution requires a single partition and the input has
+/// multiple partitions. This prevents `SortExec(preserve_partitioning=true)`
+/// from violating `SinglePartition` requirements.
+pub fn add_sort_above_with_distribution<T: Clone + Default>(
+ node: PlanContext<T>,
+ sort_requirements: LexRequirement,
+ fetch: Option<usize>,
+ required_distribution: &Distribution,
+) -> PlanContext<T> {
+ let mut sort_reqs: Vec<_> = sort_requirements.into();
+ sort_reqs.retain(|sort_expr| {
+ node.plan
+ .equivalence_properties()
+ .is_expr_constant(&sort_expr.expr)
+ .is_none()
+ });
+ let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::<Vec<_>>();
+ let Some(ordering) = LexOrdering::new(sort_exprs) else {
+ return node;
+ };
+ let input_has_multiple_partitions =
+ node.plan.output_partitioning().partition_count() > 1;
+
+ let mut new_sort =
+ SortExec::new(ordering.clone(), Arc::clone(&node.plan)).with_fetch(fetch);
+ if input_has_multiple_partitions {
+ new_sort = new_sort.with_preserve_partitioning(true);
+ }
+
+ let sort_node = PlanContext::new(Arc::new(new_sort), T::default(), vec![node]);
+
+ // If the parent requires SinglePartition and the input has multiple partitions,
+ // wrap the partition-preserving sort in SortPreservingMergeExec.
+ if matches!(required_distribution, Distribution::SinglePartition)
+ && input_has_multiple_partitions
+ {
+ PlanContext::new(
+ Arc::new(
+ SortPreservingMergeExec::new(ordering, Arc::clone(&sort_node.plan))
+ .with_fetch(fetch),
+ ),
+ T::default(),
+ vec![sort_node],
+ )
+ } else {
+ sort_node
+ }
+}
+
/// This utility function adds a `SortExec` above an operator according to the
/// given ordering requirements while preserving the original partitioning. If
/// requirement is already satisfied no `SortExec` is added.
diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt
index 67d2c1e7b5..3d0e5aac5c 100644
--- a/datafusion/sqllogictest/test_files/explain.slt
+++ b/datafusion/sqllogictest/test_files/explain.slt
@@ -237,9 +237,8 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE
physical_plan after join_selection SAME TEXT AS ABOVE
physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE
physical_plan after FilterPushdown SAME TEXT AS ABOVE
-physical_plan after EnforceDistribution SAME TEXT AS ABOVE
+physical_plan after EnsureRequirements SAME TEXT AS ABOVE
physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE
-physical_plan after EnforceSorting SAME TEXT AS ABOVE
physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE
physical_plan after WindowTopN SAME TEXT AS ABOVE
physical_plan after ProjectionPushdown SAME TEXT AS ABOVE
@@ -318,9 +317,8 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE
physical_plan after join_selection SAME TEXT AS ABOVE
physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE
physical_plan after FilterPushdown SAME TEXT AS ABOVE
-physical_plan after EnforceDistribution SAME TEXT AS ABOVE
+physical_plan after EnsureRequirements SAME TEXT AS ABOVE
physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE
-physical_plan after EnforceSorting SAME TEXT AS ABOVE
physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE
physical_plan after WindowTopN SAME TEXT AS ABOVE
physical_plan after ProjectionPushdown SAME TEXT AS ABOVE
@@ -365,9 +363,8 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE
physical_plan after join_selection SAME TEXT AS ABOVE
physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE
physical_plan after FilterPushdown SAME TEXT AS ABOVE
-physical_plan after EnforceDistribution SAME TEXT AS ABOVE
+physical_plan after EnsureRequirements SAME TEXT AS ABOVE
physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE
-physical_plan after EnforceSorting SAME TEXT AS ABOVE
physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE
physical_plan after WindowTopN SAME TEXT AS ABOVE
physical_plan after ProjectionPushdown SAME TEXT AS ABOVE
@@ -614,9 +611,8 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE
physical_plan after join_selection SAME TEXT AS ABOVE
physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE
physical_plan after FilterPushdown SAME TEXT AS ABOVE
-physical_plan after EnforceDistribution SAME TEXT AS ABOVE
+physical_plan after EnsureRequirements SAME TEXT AS ABOVE
physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE
-physical_plan after EnforceSorting SAME TEXT AS ABOVE
physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE
physical_plan after WindowTopN SAME TEXT AS ABOVE
physical_plan after ProjectionPushdown SAME TEXT AS ABOVE
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org