[EPIC] Merge EnforceDistribution + EnforceSorting into a single EnsureRequirements rule for correctness and idempotency

Summary

EnforceDistribution and EnforceSorting are currently implemented as two separate physical optimizer rules that run independently. This design is unique among major query engines -- Spark (EnsureRequirements), Presto/Trino (AddExchanges), and others all handle distribution and sorting in a single combined rule. The separation in DataFusion leads to a class of correctness bugs where one rule undoes the invariants established by the other, because sorting and distribution are fundamentally coupled (preserve_partitioning on SortExec directly affects output partition count).

This epic tracks the work to merge these two rules into a single EnsureRequirements-style rule that handles both distribution and sorting in one pass, eliminating the non-idempotent composition and the recurring bugs it causes.

Motivation

The fundamental coupling

SortExec has a preserve_partitioning flag that determines whether it outputs one partition (merging all inputs) or N partitions (sorting each independently). This means every sorting decision is also a distribution decision, and vice versa. Handling them in separate rules creates a semantic gap where each rule makes locally correct decisions that are globally incorrect.

Non-idempotent composition

Running EnforceDistribution followed by EnforceSorting does not produce a stable plan. Running the pair again can produce a different (and sometimes invalid) plan:

Round 1: EnforceDistribution -> fixes distribution
         EnforceSorting      -> pushdown_sorts breaks distribution
         
Round 2: EnforceDistribution -> fixes the NEW distribution violation
         EnforceSorting      -> pushdown_sorts breaks it AGAIN (different location)

This is particularly problematic for downstream projects that run custom optimizer rules between or after these passes (e.g., remote execution, materialized view selection), which necessitate additional rounds of EnforceSorting.

Real-world impact

We maintain a production system (Polygon.io Atlas) serving financial market data APIs built on DataFusion. The separation of these rules has caused multiple production incidents over the past months:

SanityCheckPlan failures on multi-partition StorageExec + GlobalLimitExec (April 2026): EnforceSorting's pushdown_sorts pushed a SortExec through an intermediate node onto a 32-partition input, setting preserve_partitioning=true without inserting SortPreservingMergeExec. GlobalLimitExec requires SinglePartition -> 502 errors for specific API users. Root cause: pushdown_sorts has no knowledge of distribution requirements.
InterleaveExec::with_new_children panics (InterleaveExec::with_new_children panics when optimizer rewrites change children's partitioning #21826, April 2026): Running EnforceDistribution twice (which our custom optimizer chain does) causes InterleaveExec created in the first pass to panic when the second pass changes children's partitioning.
Planning time explosion with materialized views (April 2026): EnforceDistribution's adjust_input_keys_ordering returns Transformed::yes unconditionally (EnforceDistribution: adjust_input_keys_ordering returns Transformed::yes unconditionally for non-join plans #21946), causing unnecessary plan tree rebuilds that trigger expensive cost re-evaluation in OneOfExec (materialized view candidate selection). Each EnforceSorting + EnforceDistribution round compounds this cost.

Existing upstream issues (same root cause)

Bug: applying multiple times EnforceDistribution generates invalid plan #14150: "Bug: applying multiple times EnforceDistribution generates invalid plan" -- running EnforceDistribution twice loses limit fetch values, producing wrong results.
Sanity check failed when sort and aggregate on a multi-partitioned table #18989: "Sanity check failed when sort and aggregate on a multi-partitioned table" -- EnforceDistribution fails to inject necessary RepartitionExec between aggregate operations.
Window aggregates output order broken due to hash repartitioning #16888: "Window aggregates output order broken due to hash repartitioning" -- EnforceDistribution breaks ordering established by EnforceSorting.
InterleaveExec::with_new_children panics when optimizer rewrites change children's partitioning #21826: "InterleaveExec::with_new_children panics when optimizer rewrites change children's partitioning" -- second EnforceDistribution pass panics on plan created by first pass.

All of these share the same root cause: two separate rules making independent decisions about coupled concerns.

How other engines solve this

Apache Spark: `EnsureRequirements`

Spark handles both in a single rule using transformUp:

// For each operator, in a single pass:
// 1. Check requiredChildDistribution -> add ShuffleExchangeExec if needed
// 2. Check requiredChildOrdering -> add SortExec if needed
// Distribution is always resolved before sorting for the same operator.

Key design: sorting decisions are made AFTER distribution is finalized for each operator. SortExec(global=false) preserves partition boundaries. No separate "sort pushdown" pass exists.

Presto/Trino: `AddExchanges`

Presto's AddExchanges rule similarly handles both distribution and sorting properties in a single rule, using a PreferredProperties structure that carries both distribution and ordering preferences through the plan tree.

Proposed approach

Phase 1: Make `pushdown_sorts` distribution-aware (short-term fix)

Add a distribution_requirement field to ParentRequirements in pushdown_sorts, so that add_sort_above knows when to insert SortPreservingMergeExec. This is a targeted fix for the most critical bug.

Status: Implemented in our internal fork, ready to port upstream as part of this epic.

Phase 2: Add idempotency tests (validation)

Add tests that verify EnforceDistribution -> EnforceSorting produces a stable plan:

fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
    let p1 = EnforceDistribution::optimize(plan)?;
    let p1 = EnforceSorting::optimize(p1)?;
    let p2 = EnforceDistribution::optimize(p1.clone())?;
    let p2 = EnforceSorting::optimize(p2)?;
    assert_eq!(display(p1), display(p2));
}

Test against various plan topologies: multi-partition sorts with limits, unions with mixed partition counts, projections over multi-partition sources, window functions, etc.

Phase 3: Merge into `EnsureRequirements` (architectural fix)

Create a new EnsureRequirements rule that replaces both EnforceDistribution and EnforceSorting:

pub struct EnsureRequirements;

impl PhysicalOptimizerRule for EnsureRequirements {
    fn optimize(&self, plan: Arc<dyn ExecutionPlan>, config: &ConfigOptions) -> Result<Arc<dyn ExecutionPlan>> {
        // Single bottom-up pass:
        // For each operator, ensure children satisfy both
        // requiredChildDistribution AND requiredChildOrdering.
        // Distribution is resolved before ordering for each operator.
        plan.transform_up(|node| ensure_requirements(node, config))
    }
}

struct Requirements {
    distribution: Distribution,
    ordering: Option<OrderingRequirements>,
    fetch: Option<usize>,
}

fn ensure_requirements(node: PlanContext<Requirements>) -> Result<Transformed<...>> {
    for (child, required_dist, required_ordering) in zip(children, distributions, orderings) {
        // Step 1: Ensure distribution
        if !child.output_partitioning().satisfies(&required_dist) {
            child = add_exchange(child, required_dist);
        }
        // Step 2: Ensure ordering (distribution is already correct)
        if !child.output_ordering().satisfies(&required_ordering) {
            child = add_sort(child, required_ordering, required_dist); // dist-aware!
        }
    }
}

Key properties:

Single pass: No separate pushdown_sorts that can undo distribution work
Distribution before sorting: For each operator, distribution is settled before sorting decisions, like Spark
Naturally idempotent: Running it twice produces the same plan because each operator's children are already correct after the first pass
Sort pushdown integrated: Instead of a separate top-down pass, sort pushdown is handled by the bottom-up pass recognizing when a child already satisfies ordering (no sort needed)

Migration path

EnsureRequirements can coexist with the old rules during development
Add a feature flag to switch between old and new behavior
Validate with the full DataFusion test suite + sqllogictest
Deprecate EnforceDistribution + EnforceSorting after stabilization

Sub-tasks

References

Spark EnsureRequirements: https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
Presto AddExchanges: https://github.com/prestodb/presto/wiki/New-Optimizer
Bug: applying multiple times EnforceDistribution generates invalid plan #14150: Multiple EnforceDistribution generates invalid plan
Sanity check failed when sort and aggregate on a multi-partitioned table #18989: Sanity check failed with sort + aggregate on multi-partitioned table
Window aggregates output order broken due to hash repartitioning #16888: Window aggregates output order broken due to hash repartitioning
InterleaveExec::with_new_children panics when optimizer rewrites change children's partitioning #21826: InterleaveExec::with_new_children panics from optimizer rewrites
EnforceDistribution: adjust_input_keys_ordering returns Transformed::yes unconditionally for non-join plans #21946: adjust_input_keys_ordering returns Transformed::yes unconditionally

Initial Design Thoughts

After analyzing both rules in depth (~4000 lines total), here's my proposed architecture for the merged EnsureRequirements rule.

Why merging is the right approach

The current separation has a fundamental design flaw: sorting and distribution are coupled through SortExec.preserve_partitioning, but the two rules make decisions independently. This coupling means:

EnforceDistribution already partially handles ordering (replace_order_preserving_variants, add_sort_above_with_check)
EnforceSorting already partially handles distribution (parallelize_sorts removes CoalescePartitionsExec, analyze_immediate_sort_removal adds SortPreservingMergeExec)
The responsibility boundary is blurred — both rules touch both concerns

Every major query engine (Spark's EnsureRequirements, Presto/Trino's AddExchanges) handles distribution and sorting in a single rule for exactly this reason.

The `pushdown_sorts` problem

The root cause of non-idempotency is pushdown_sorts — the only top-down pass in EnforceSorting. It:

Removes existing SortExec (which ensure_sorting just placed correctly)
Walks down the tree carrying ordering requirements
Reinserts SortExec at the deepest pushable position via add_sort_above
add_sort_above has no distribution context → sets preserve_partitioning=true on multi-partition input without inserting SortPreservingMergeExec

Spark doesn't have a separate pushdown_sorts pass. Sort pushdown is implicit in the bottom-up traversal: if a child already satisfies the ordering requirement, no SortExec is added — equivalent to "the sort was pushed to the deepest satisfying position."

Proposed architecture

EnsureRequirements::optimize(plan)
│
├─ Phase 1 (optional): reorder_join_keys (top-down)
│   └─ Keep existing adjust_input_keys_ordering logic unchanged
│
├─ Phase 2: ensure_requirements (single bottom-up pass)
│   └─ For each node (bottom-up):
│       For each child:
│         Step 1: Ensure distribution requirement
│           └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
│         Step 2: Ensure ordering requirement (distribution-aware)
│           └─ Add SortExec with correct preserve_partitioning + SPM if needed
│         Step 3: Local optimizations
│           ├─ parallelize_sorts (inline): Coalesce+Sort → Sort+SPM
│           ├─ order-preserving variants (inline): Repartition → preserve_order
│           └─ partial sort (inline): SortExec → PartialSortExec
│
└─ No separate pushdown_sorts pass!
    └─ Sort pushdown is implicit: bottom-up pass only adds sort
       where child doesn't satisfy ordering → sort is naturally
       at the deepest satisfying position

Unified data structure

/// Combined requirement for a single child
struct ChildRequirement {
    distribution: Distribution,
    ordering: Option<OrderingRequirements>,
    fetch: Option<usize>,
    // Auxiliary decision flags
    hash_necessary: bool,
    roundrobin_beneficial: bool,
}

Key design properties

Single pass: One bottom-up traversal handles both distribution and sorting
Distribution before sorting: For each operator's children, distribution is settled first, then sorting decisions are made with full distribution context
Naturally idempotent: Running the rule twice produces the same plan — each child is already correct after the first pass
No OutputRequirementExec needed: The single pass carries requirements directly, eliminating the add/remove ceremony and the require_top_ordering_helper distribution extraction bug
pushdown_sorts eliminated: Sort pushdown is implicit in the bottom-up traversal

Migration strategy

Implement EnsureRequirements alongside existing rules
Feature flag to switch between old and new behavior
Validate against full DataFusion test suite + sqllogictest
Deprecate EnforceDistribution + EnforceSorting after stabilization

Scope & complexity

This is a large change (~4000 lines to rewrite), but the existing code already has significant overlap between the two rules. The merged version should be significantly simpler because:

No need for OutputRequirementExec add/remove ceremony
No pushdown_sorts top-down pass
No duplicated replace_order_preserving_variants logic
Single traversal instead of 5+ traversals

I'll start with Phase 1 (the pushdown_sorts distribution fix) as a standalone upstream PR, then iterate on the full merge.

Testing Strategy — No Regressions Allowed

The most critical requirement for this architectural change: zero test regressions + comprehensive new coverage for all known production bugs.

Existing test coverage that MUST pass

All existing enforce_sorting tests (datafusion/core/tests/physical_optimizer/enforce_sorting.rs) — 160+ tests
All existing enforce_distribution tests (datafusion/core/tests/physical_optimizer/enforce_distribution.rs) — 200+ tests
All sqllogictest tests — the full SLT suite must pass unchanged
SanityCheckPlan tests — all sanity checks must continue to pass

New tests to add (from real production incidents)

These are bugs we've hit in production that the merged rule must prevent:

Multi-partition sort + limit + SanityCheckPlan (the bug that triggered this epic):
- OutputRequirementExec(SinglePartition) → StorageExec(N partitions)
- EnforceSorting must insert SortPreservingMergeExec between SortExec(preserve_partitioning=true) and GlobalLimitExec
- Test with 2, 16, and 32 partitions
Sort pushdown through projection over multi-partition input:
- OutputRequirementExec(SinglePartition) → ProjectionExec → RepartitionExec(10)
- Sort pushed through projection must include SortPreservingMergeExec
Sort pushdown with already-optimal plan (no extra SPM):
- OutputRequirementExec(SinglePartition) → SortPreservingMergeExec → SortExec(preserve=true) → RepartitionExec(10)
- Must NOT add duplicate SortPreservingMergeExec
InterleaveExec stability under multiple passes (InterleaveExec::with_new_children panics when optimizer rewrites change children's partitioning #21826):
- UnionExec children with varying partitioning
- Running optimizer twice must not panic InterleaveExec::with_new_children
Multiple EnforceDistribution preserves fetch (Bug: applying multiple times EnforceDistribution generates invalid plan #14150):
- SELECT * FROM t ORDER BY c LIMIT 5 with multi-partition input
- Fetch value must survive multiple optimizer passes
Sort + aggregate on multi-partitioned table (Sanity check failed when sort and aggregate on a multi-partitioned table #18989):
- Sort → Aggregate → Sort → Aggregate on multi-partitioned input
- Must insert proper RepartitionExec between aggregates

Idempotency tests (NEW — the core guarantee)

/// Framework: run EnsureRequirements twice, assert same plan
fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
    let config = ConfigOptions::default();
    let p1 = EnsureRequirements::new().optimize(plan, &config).unwrap();
    let p2 = EnsureRequirements::new().optimize(p1.clone(), &config).unwrap();
    assert_eq!(
        displayable(p1.as_ref()).indent(true).to_string(),
        displayable(p2.as_ref()).indent(true).to_string(),
        "EnsureRequirements is not idempotent"
    );
    // Also verify SanityCheckPlan passes on both
    SanityCheckPlan::new().optimize(p1, &config).unwrap();
    SanityCheckPlan::new().optimize(p2, &config).unwrap();
}

Test topologies for idempotency:

Multi-partition sort + limit
Union with mixed partition counts
Projection over multi-partition source
Window functions over repartitioned data
Hash join with sort-merge fallback
Aggregate with grouping sets

Test-first approach

Each sub-task in this epic should follow:

Write the failing test first (from production incident or known bug)
Implement the fix
Verify no existing tests regress
Add idempotency test for the new topology

This ensures the merged rule is strictly better — it passes everything the old rules passed, plus catches everything they missed.

Add experimental EnsureRequirementsNoPushdown variant (Phase 3 WIP)

Fix EnforceDistribution idempotency: preserve fetch across passes (ap…

Replace default EnforceDistribution + EnforceSorting with EnsureRequi…

Add EnsureRequirements: idempotent merged EnforceDistribution + Enfor…

No pushdown_sorts top-down pass

I have a question about how to eliminate pushdown_sorts. For example:

case 1:
sort(k1 asc)
-- filter (k1 > 0)
---- scan(existing file order: k1 asc)

case 2:
sort(k1 desc)
-- filter (k1 > 0)
---- scan(existing file order: k1 asc)  // file format supports reverse scan

For case 1, a bottom-up traversal can eliminate the SortExec, because the existing input order of SortExec already satisfies the requirement. For case 2, however, it seems that some top-down negotiation is still needed: SortExec has to first tell the scanner to perform a reverse scan, and then the bottom-up traversal can remove the redundant sort.

I was thinking that maybe we can let the EnsureRequirements rule only do a bottom-up scan and insert all required operators for ordering and repartitioning requirements. Other decisions, such as exploring reversed file scan order in the example above, can be kept in other optimizer rules. This would make EnsureRequirements easier to maintain.

Another question: the current EnforceSorting rule seems to handle a lot of special cases:

datafusion/datafusion/physical-optimizer/src/enforce_sorting/mod.rs

Line 198 in e62f06c

/// Optimizer consists of 5 main parts which work sequentially

I'm not sure how all those special cases are caused. Perhaps a unified bottom-up approach with EnsureRequirements can eliminate some of them. We might also want to separate some existing edge cases out to keep the EnsureRequirements rule simpler and more focused.

Here is one observation that might simplify the implementation a bit:

Currently, in physical planning, SortExec is ambiguous. In the initial physical plan, the logical plan's sort node is directly converted to SortExec. In this case, it is more like a placeholder that is actually a "logical" sort operator, and it relies on later optimizer passes to add the required SortPreservingMergeExec to make the whole plan valid.

Meanwhile, for some operators that have input requirements, specifically SortMergeJoin and BoundedWindowAggExec, their initial physical plans don't include SortExec. Instead, they directly use ExecutionPlan::required_input_ordering to let later optimizer passes insert the required sort operators.

So we might choose either approach and make them consistent. This could make the later implementation require fewer special cases. For example, we could convert the logical plan's sort operator to OutputRequirementExec in the initial physical plan, so sort requirements have a consistent representation. Alternatively, we could insert SortExec below SortMergeJoin / BoundedWindowAggExec if that approach is easier.

Add EnsureRequirements: idempotent merged EnforceDistribution + Enfor…

docs: tighten EnsureRequirements module doc; add explicit apache#21973 …

Add EnsureRequirements: idempotent merged EnforceDistribution + Enfor…

docs: tighten EnsureRequirements module doc; add explicit apache#21973 …

@zhuqi-lucas

@zhuqi-lucas is this epic done now that we merged #21976 ?

@alamb

@zhuqi-lucas is this epic done now that we merged #21976 ?

Yes, thanks @alamb , close this epic now.

Navigation Menu

Repository navigation

[EPIC] Merge EnforceDistribution + EnforceSorting into a single EnsureRequirements rule for correctness and idempotency #21973

Summary

Motivation

The fundamental coupling

Non-idempotent composition

Real-world impact

Existing upstream issues (same root cause)

How other engines solve this

Apache Spark: `EnsureRequirements`

Presto/Trino: `AddExchanges`

Proposed approach

Phase 1: Make `pushdown_sorts` distribution-aware (short-term fix)

Phase 2: Add idempotency tests (validation)

Phase 3: Merge into `EnsureRequirements` (architectural fix)

Migration path

Sub-tasks

References

Initial Design Thoughts

Why merging is the right approach

The `pushdown_sorts` problem

Proposed architecture

Unified data structure

Key design properties

Migration strategy

Scope & complexity

Testing Strategy — No Regressions Allowed

Existing test coverage that MUST pass

New tests to add (from real production incidents)

Idempotency tests (NEW — the core guarantee)

Test-first approach

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[EPIC] Merge EnforceDistribution + EnforceSorting into a single EnsureRequirements rule for correctness and idempotency #21973

Description

Summary

Motivation

The fundamental coupling

Non-idempotent composition

Real-world impact

Existing upstream issues (same root cause)

How other engines solve this

Apache Spark: EnsureRequirements

Presto/Trino: AddExchanges

Proposed approach

Phase 1: Make pushdown_sorts distribution-aware (short-term fix)

Phase 2: Add idempotency tests (validation)

Phase 3: Merge into EnsureRequirements (architectural fix)

Migration path

Sub-tasks

References

Activity

zhuqi-lucas commented on May 1, 2026

Initial Design Thoughts

Why merging is the right approach

The pushdown_sorts problem

Proposed architecture

Unified data structure

Key design properties

Migration strategy

Scope & complexity

zhuqi-lucas commented on May 1, 2026

Testing Strategy — No Regressions Allowed

Existing test coverage that MUST pass

New tests to add (from real production incidents)

Idempotency tests (NEW — the core guarantee)

Test-first approach

2010YOUY01 commented on May 10, 2026

alamb commented on Jun 25, 2026

zhuqi-lucas commented on Aug 9, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Apache Spark: `EnsureRequirements`

Presto/Trino: `AddExchanges`

Phase 1: Make `pushdown_sorts` distribution-aware (short-term fix)

Phase 3: Merge into `EnsureRequirements` (architectural fix)

The `pushdown_sorts` problem