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

[zhuqi-lucas]

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:

1. 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.

2. 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.

3. 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

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

Sub-tasks

• Phase 1: Port pushdown_sorts distribution fix upstream (will be submitted as upstream PR)
• Phase 1: Port ensure_sorting distribution fix upstream (will be submitted as upstream PR)
• Phase 2: Add idempotency test framework for physical optimizer rules
• Phase 2: Add idempotency tests for EnforceDistribution + EnforceSorting composition
• Phase 3: Design EnsureRequirements API and Requirements structure
• Phase 3: Implement combined distribution + ordering enforcement in single bottom-up pass
• Phase 3: Integrate sort pushdown into the bottom-up pass
• Phase 3: Handle parallelize_sorts and replace_with_order_preserving_variants optimizations
• Phase 3: Migrate OutputRequirements (add/remove) into the new rule
• Phase 3: Feature flag and migration path
• Phase 3: Validate against full test suite

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

[zhuqi-lucas]

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:

1. EnforceDistribution already partially handles ordering (replace_order_preserving_variants, add_sort_above_with_check)
2. EnforceSorting already partially handles distribution (parallelize_sorts removes CoalescePartitionsExec, analyze_immediate_sort_removal adds SortPreservingMergeExec)
3. 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:

1. Removes existing SortExec (which ensure_sorting just placed correctly)
2. Walks down the tree carrying ordering requirements
3. Reinserts SortExec at the deepest pushable position via add_sort_above
4. 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

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

Migration strategy

1. Implement EnsureRequirements alongside existing rules
2. Feature flag to switch between old and new behavior
3. Validate against full DataFusion test suite + sqllogictest
4. 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.

[zhuqi-lucas]

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

1. All existing enforce_sorting tests (datafusion/core/tests/physical_optimizer/enforce_sorting.rs) — 160+ tests
2. All existing enforce_distribution tests (datafusion/core/tests/physical_optimizer/enforce_distribution.rs) — 200+ tests
3. All sqllogictest tests — the full SLT suite must pass unchanged
4. 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:

1. 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

2. Sort pushdown through projection over multi-partition input:

   • OutputRequirementExec(SinglePartition) → ProjectionExec → RepartitionExec(10)
   • Sort pushed through projection must include SortPreservingMergeExec

3. Sort pushdown with already-optimal plan (no extra SPM):

   • OutputRequirementExec(SinglePartition) → SortPreservingMergeExec → SortExec(preserve=true) → RepartitionExec(10)
   • Must NOT add duplicate SortPreservingMergeExec

4. 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

5. 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

6. 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:

1. Write the failing test first (from production incident or known bug)
2. Implement the fix
3. Verify no existing tests regress
4. 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.