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/24 02:57:49 UTC
(datafusion) branch main updated: Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts (#21976)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new fb1c0f342f Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts (#21976)
fb1c0f342f is described below
commit fb1c0f342f495f502e1a9002a74e2fb038919191
Author: Qi Zhu <qi...@polygon.io>
AuthorDate: Sun May 24 10:37:08 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 a6cddf200a..0df26c4274 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