Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts

[zhuqi-lucas]

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 (Bug: applying multiple times EnforceDistribution generates invalid plan #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 (#21976 (comment)), 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 Bug: applying multiple times EnforceDistribution generates invalid plan #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).

[Copilot]

Pull request overview

This PR replaces the separate physical optimizer rules EnforceDistribution and EnforceSorting with a unified EnsureRequirements rule, and updates sort pushdown to be distribution-aware to prevent invalid plans when SinglePartition is required (e.g., under GlobalLimitExec). It also addresses fetch preservation when EnforceDistribution strips and re-adds partition-changing operators, and updates reference documentation / SLT expectations accordingly.

Changes:

• Introduces EnsureRequirements and wires it into the default physical optimizer chain in place of EnforceDistribution + EnforceSorting.
• Makes pushdown_sorts distribution-aware by propagating distribution requirements and inserting SortPreservingMergeExec when needed.
• Preserves fetch across EnforceDistribution’s removal/reinsertion of SortPreservingMergeExec / CoalescePartitionsExec, and updates SLT + optimizer rule docs.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
datafusion/sqllogictest/test_files/explain.slt	Updates EXPLAIN rule-name expectations to EnsureRequirements.
datafusion/physical-optimizer/src/utils.rs	Adds add_sort_above_with_distribution to insert SPM when parent requires SinglePartition.
datafusion/physical-optimizer/src/optimizer.rs	Swaps default optimizer chain to use EnsureRequirements instead of separate distribution/sorting rules.
datafusion/physical-optimizer/src/lib.rs	Exposes new ensure_requirements module.
datafusion/physical-optimizer/src/ensure_requirements/new_tests.rs	Adds extensive tests for correctness/idempotency across topologies.
datafusion/physical-optimizer/src/ensure_requirements/mod.rs	Implements EnsureRequirements (and an experimental no-pushdown variant) plus additional tests/docs.
datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs	Propagates distribution requirements through sort pushdown; uses distribution-aware sort insertion.
datafusion/physical-optimizer/src/enforce_sorting/mod.rs	Makes replace_with_partial_sort public for reuse.
datafusion/physical-optimizer/src/enforce_distribution.rs	Preserves fetch when stripping SPM/Coalesce and reapplies it when re-merging.
datafusion/core/src/optimizer_rule_reference.md	Updates docs to reflect EnsureRequirements in the default physical pipeline.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[zhuqi-lucas]

cc @alamb @adriangb @xudong963 @Dandandan @2010YOUY01

We've hit several production bugs caused by EnforceDistribution and EnforceSorting running as separate rules — they're coupled through SortExec.preserve_partitioning, which makes the composition non-idempotent.
Other engines (Spark's EnsureRequirements, Trino's AddExchanges) handle both in a single rule, which sidesteps this entire bug class.

I initially attempted a true single-pass rewrite, but the interactions between distribution and sorting requirements turned out to be deeper than anticipated.

This PR takes a more conservative approach — composing the existing rules into one with the necessary correctness fixes. A clean single-pass design is left as a follow-up once the foundation is in place.

If a default change feels too risky for a first step, I'm happy to gate it behind a config flag (e.g. datafusion.optimizer.use_ensure_requirements) so users can opt in and we can validate equivalence in CI before flipping the default.

Would appreciate any thoughts on the overall direction.

[2010YOUY01]

Thank you! I think this is the right direction. The existing two rules both
describe the expected input shape of a plan, so they should be tightly coupled.
I believe this direction can reduce entropy in the long run.

I do not fully understand the details yet, but I will try to spend more time looking
into them and continue the review.

[2010YOUY01]

Suggested change

	// EnsureRequirements: merged EnforceDistribution + EnforceSorting into a
	// single idempotent rule with distribution-aware pushdown_sorts.
	// See https://github.com/apache/datafusion/issues/21973
	// 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 with
	// distribution-aware `pushdown_sorts`.
	// See https://github.com/apache/datafusion/issues/21973.

Added more comments since this is the entry point.

I also have a question: What is this 'distribution-aware pushdown_sorts'? We cloud link some reference to it.

[zhuqi-lucas]

Thanks @2010YOUY01 for review and good suggestion!

Will apply the suggestion in the next push (per alamb's plan I'll land a test-extraction PR first and rebase this on top).

On "distribution-aware pushdown_sorts": the old pushdown_sorts didn't propagate the parent's required_input_distribution, so a SortExec could get pushed below something requiring SinglePartition (e.g. under GlobalLimitExec) — giving an invalid plan because the sort then runs on multiple partitions. The fix propagates the requirement and inserts a SortPreservingMergeExec where needed. Will link to #21973 in the comment block when I update.

[alamb]

Thanks @zhuqi-lucas -- this is exciting and a really nice idea.

Given how important this infrastructure is for all DataFusion plans I think the key to making this change will be to make sure we don't introduce regressions.

My biggest concern is that I can't really evaluate what, if any, tests were changed / added / removed in this PR (though it is impressive that all slt is passing)

Initial thoughts

1. The existing EnforceDistribtuon and EnforceSorting code should be removed
2. There are a LOT of existing unit tests for EnforceDistribtuon and EnforceSorting -- it looks like most of these have been ported to the new single pass, but it is really hard for me to verify

Suggested next steps

As there are already several large changes queued up in the 54 release, Maybe we can work to get the tests into shape and plan to merge this PR after we ship 54.

What I suggest we do is:

1. Prepare the tests
   First, make a separate PR to move all existing tests for enforce sorting and enforce distribution somewhere outside the existing modules. By having the tests in a separate module, it will be much easier to understand what, if anything, changed when you replace EnforceSorting/EnforceDistribtuion with EnforceRequirements

This first PR should be fast/easy to review

2. Then rebase this PR

Given the tests are separate from the code, it will then be much easier to see / verify that this unified pass both: 1) has the same test coverage, and 2) keeps the same behavior

[alamb]

Why is this a separate module? If we want tests in a separate module, I think putting it in an integration test might be easier to find (datafusion/physical-optimizer/tests/ensure_requirements.rs perhaps?)

[alamb]

many of these tests look similar/ the same as what is in enforce distribution, etc but it is hard to tell what changed

[zhuqi-lucas]

Thanks @alamb. Reading again I see two separate asks:

• The inline on new_tests.rs (move it to datafusion/physical-optimizer/tests/ensure_requirements.rs) is something I can do inside this PR right away — will push that next.
• The bigger one (move the existing core/tests/physical_optimizer/{enforce_distribution.rs, enforce_sorting.rs, ...} integration tests, ~9.8k lines, to live alongside the optimizer crate) — I'll do that as a separate prep PR after 54 ships, then rebase this on top so the porting diff is clean.

Will ping you when the inline fix lands and again when the prep PR is up.

[xudong963]

What I suggest we do is:

1. Prepare the tests
   First, make a separate PR to move all existing tests for enforce sorting and enforce distribution somewhere outside the existing modules. By having the tests in a separate module, it will be much easier to understand what, if anything, changed when you replace EnforceSorting/EnforceDistribtuion with EnforceRequirements

+1 for this

[zhuqi-lucas]

What I suggest we do is:

1. Prepare the tests
   First, make a separate PR to move all existing tests for enforce sorting and enforce distribution somewhere outside the existing modules. By having the tests in a separate module, it will be much easier to understand what, if anything, changed when you replace EnforceSorting/EnforceDistribtuion with EnforceRequirements

+1 for this

Thanks @xudong963 , will go this way.

[alamb]

Thanks @alamb. Reading again I see two separate asks:

• The inline on new_tests.rs (move it to datafusion/physical-optimizer/tests/ensure_requirements.rs) is something I can do inside this PR right away — will push that next.
• The bigger one (move the existing core/tests/physical_optimizer/{enforce_distribution.rs, enforce_sorting.rs, ...} integration tests, ~9.8k lines, to live alongside the optimizer crate) — I'll do that as a separate prep PR after 54 ships, then rebase this on top so the porting diff is clean.

Will ping you when the inline fix lands and again when the prep PR is up.

Thanks @zhuqi-lucas

FYI I think we can reorganize the tests before 54 as that has a low chance of introducing bugs (so doesn't need as long a "bake" time)

In terms of this PR, my plan for reviewing it is basically

1. Review what, if any tests have changed (will spend most effort here)
2. Run sql planning performance benchmarks to ensure no regressions
3. Review the code for anything that stands out

So TLDR is I plan to spend most of my time doing verification rather than code review

[zhuqi-lucas]

Thanks @alamb. Reading again I see two separate asks:

• The inline on new_tests.rs (move it to datafusion/physical-optimizer/tests/ensure_requirements.rs) is something I can do inside this PR right away — will push that next.
• The bigger one (move the existing core/tests/physical_optimizer/{enforce_distribution.rs, enforce_sorting.rs, ...} integration tests, ~9.8k lines, to live alongside the optimizer crate) — I'll do that as a separate prep PR after 54 ships, then rebase this on top so the porting diff is clean.

Will ping you when the inline fix lands and again when the prep PR is up.

Thanks @zhuqi-lucas

FYI I think we can reorganize the tests before 54 as that has a low chance of introducing bugs (so doesn't need as long a "bake" time)

In terms of this PR, my plan for reviewing it is basically

1. Review what, if any tests have changed (will spend most effort here)
2. Run sql planning performance benchmarks to ensure no regressions
3. Review the code for anything that stands out

So TLDR is I plan to spend most of my time doing verification rather than code review

Thanks @alamb !

Test-extraction PR is up: #22117. Hit a snag with the original "move to physical-optimizer/tests/" approach — the tests use SessionContext so they'd need datafusion as a dev-dep of datafusion-physical-optimizer, which dev/depcheck rejects (it'd block crates.io publish through the resulting cycle).

Pivoted to a new workspace crate datafusion-physical-optimizer-tests (publish = false, following the precedent of test-utils and datafusion-wasmtest). The crate is a leaf in the dependency graph — nothing depends on it — so its dev-deps on datafusion and datafusion-physical-optimizer introduce no cycle. cargo metadata confirms zero incoming edges.

Git detects all four moves as 93–98% renames. Once #22117 lands I'll rebase this PR on top, and the resulting diff should be the function-name swap from EnforceDistribution::new() / EnforceSorting::new() to EnsureRequirements::new() on the moved files plus the genuinely new tests — which is the comparison you wanted to be able to do.

Happy to discuss the new-crate choice on #22117 if you'd prefer a different structure.

[alamb]

Follow on comment

• test: extract enforce_distribution / enforce_sorting tests into a dedicated crate #22117 (comment)

[zhuqi-lucas]

Thanks @alamb, that direction worked out really cleanly. All three steps are in now.

The old EnforceDistribution / EnforceSorting rule structs are gone. The existing tests in core/tests/physical_optimizer/ now call EnsureRequirements::new() directly — kept the original test framework shape so the historical [Dist, Sort] / [Sort, Dist, Sort] sequences are still there, they just collapse to running the merged rule a couple of times (and the merged rule is idempotent, which is the whole point).

Snapshots refreshed; the consistent pattern is SortExec + CoalescePartitionsExec → SortPreservingMergeExec. EnsureRequirements is now also pulling in parallelize_sorts + replace_with_order_preserving_variants on plan shapes that the bare single-rule path used to miss, so a chunk of the diff is the optimizer producing better plans on the same inputs. Worth a careful look in review.

On new_tests.rs: looking at it more carefully, every test in that file (and every inline test in ensure_requirements/mod.rs) is using either an optimize_and_sanity_check helper (regression for the SanityCheckPlan validation failures that motivated this PR) or assert_idempotent — none of them are smoke tests, they're all exercising the new properties. The real redundancy was that the two files duplicated the helpers and a lot of plan shapes. Dropped new_tests.rs and kept the mod.rs inline ones, since the named #14150 regression lives there and the plan-shape coverage (hash join, aggregate, window, projection, sort+limit, fetch preservation, …) is comprehensive already. If you'd rather I cull harder, easy follow-up.

Hit one CI failure on the hash-collisions job after the first push: EnforceSortingTest used ConfigOptions::new() (defaulted target_partitions to CPU count) and the snapshots had my local 12 baked in, so the 4-core runner blew up. Pinned target_partitions = 10 to match the convention in enforce_distribution.rs and re-accepted. Should be green now.

[github-actions]

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details

     Cloning apache/main
    Building datafusion v53.1.0 (current)
       Built [  82.374s] (current)
     Parsing datafusion v53.1.0 (current)
      Parsed [   0.034s] (current)
    Building datafusion v53.1.0 (baseline)
       Built [  81.888s] (baseline)
     Parsing datafusion v53.1.0 (baseline)
      Parsed [   0.034s] (baseline)
    Checking datafusion v53.1.0 -> v53.1.0 (no change; assume patch)
     Checked [   0.677s] 222 checks: 222 pass, 30 skip
     Summary no semver update required
    Finished [ 166.450s] datafusion
    Building datafusion-physical-optimizer v53.1.0 (current)
       Built [  36.841s] (current)
     Parsing datafusion-physical-optimizer v53.1.0 (current)
      Parsed [   0.022s] (current)
    Building datafusion-physical-optimizer v53.1.0 (baseline)
       Built [  37.171s] (baseline)
     Parsing datafusion-physical-optimizer v53.1.0 (baseline)
      Parsed [   0.022s] (baseline)
    Checking datafusion-physical-optimizer v53.1.0 -> v53.1.0 (no change; assume patch)
     Checked [   0.104s] 222 checks: 221 pass, 1 fail, 0 warn, 30 skip

--- failure struct_missing: pub struct removed or renamed ---

Description:
A publicly-visible struct cannot be imported by its prior path. A `pub use` may have been removed, or the struct itself may have been renamed or removed entirely.
        ref: https://doc.rust-lang.org/cargo/reference/semver.html#item-remove
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.47.0/src/lints/struct_missing.ron

Failed in:
  struct datafusion_physical_optimizer::enforce_sorting::EnforceSorting, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/da1bc44dbfe2ea5ef369389ef86956738e6352d0/datafusion/physical-optimizer/src/enforce_sorting/mod.rs:79
  struct datafusion_physical_optimizer::enforce_distribution::EnforceDistribution, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/da1bc44dbfe2ea5ef369389ef86956738e6352d0/datafusion/physical-optimizer/src/enforce_distribution.rs:184

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [  75.704s] datafusion-physical-optimizer
    Building datafusion-sqllogictest v53.1.0 (current)
       Built [ 134.912s] (current)
     Parsing datafusion-sqllogictest v53.1.0 (current)
      Parsed [   0.022s] (current)
    Building datafusion-sqllogictest v53.1.0 (baseline)
       Built [ 134.337s] (baseline)
     Parsing datafusion-sqllogictest v53.1.0 (baseline)
      Parsed [   0.023s] (baseline)
    Checking datafusion-sqllogictest v53.1.0 -> v53.1.0 (no change; assume patch)
     Checked [   0.089s] 222 checks: 222 pass, 30 skip
     Summary no semver update required
    Finished [ 272.864s] datafusion-sqllogictest