Skip to content

fix: normalize InterleaveExec to UnionExec before enforcing distribution - #24959

Merged
jayzhan211 merged 1 commit into
apache:mainfrom
jayzhan211:enforce-dist
Sep 6, 2026
Merged

jayzhan211 merged 1 commit into
apache:mainfrom
jayzhan211:enforce-dist

Conversation

@jayzhan211

@jayzhan211 jayzhan211 commented Sep 6, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

EnsureRequirements fails with an internal error when the plan it is given
already contains an InterleaveExec whose children no longer share a hash
(or range) partitioning:

Internal error: Assertion failed: can_interleave(children.iter()):
Can not create InterleaveExec: new children can not be interleaved.

InterleaveExec declares no distribution requirement of its own. It is only
valid while its children happen to share the same partitioning, which is why
ensure_distribution creates it from a UnionExec in the first place. Once
it exists, several rewrites can take that partitioning away: this rule strips
RepartitionExecs that nothing requires, join key reordering changes the hash
expressions, JoinSelection side swaps change output partitioning, and so on.
PlanContext::with_new_children rebuilds every parent from its updated
children while walking the tree, so the interleave is rebuilt over the changed
children before the rule's closure sees the node, and the rebuild fails.

This affects any pipeline that runs the rule more than once, re-optimizes a
deserialized plan, or runs EnforceDistribution after other rewrites (the
setup reported in #21826). It also breaks the idempotency the rule advertises.
The discussion on #21827 concluded that the fix belongs at the
EnforceDistribution call site rather than as a silent fallback in
InterleaveExec::with_new_children; this PR is that fix.

What changes are included in this PR?

  • New replace_interleave_with_union helper in
    ensure_requirements/enforce_distribution.rs, which turns an
    InterleaveExec back into the equivalent UnionExec.
  • EnsureRequirements::optimize runs it as a top-down "Phase 0" before join
    key reordering. Phase 2 then re-derives interleaves from unions wherever the
    children still qualify, exactly as it already re-derives RepartitionExec,
    CoalescePartitionsExec and SortPreservingMergeExec.
  • The pass is top-down on purpose: demoting a nested interleave first would
    change its output partitioning and make the parent interleave fail to
    rebuild.
  • Module docs updated with the new phase.

A minimal reproducer that fails on main and passes here:

// Interleave over two hash repartitions that nothing above requires.
let plan = InterleaveExec::try_new(vec![
    RepartitionExec::try_new(parquet_exec(), Partitioning::Hash(vec![col_a], 10))?,
    RepartitionExec::try_new(parquet_exec(), Partitioning::Hash(vec![col_a], 10))?,
])?;
EnsureRequirements::new().optimize(Arc::new(plan), &config)?;

What is the testing strategy for this PR?

Three new tests in
datafusion/core/tests/physical_optimizer/enforce_distribution.rs, each run
through the existing multi-pass harness (DISTRIB_DISTRIB_SORT and
SORT_DISTRIB_DISTRIB) so idempotency is checked as well:

  • interleave_falls_back_to_union_when_children_lose_partitioning: the
    reproducer above; the result is a UnionExec over the scans.
  • interleave_fallback_still_satisfies_parent_hash_requirement: same input
    under a FinalPartitioned aggregate; the aggregate gets a single hash
    repartition above the union.
  • existing_interleave_is_kept_when_children_stay_interleavable: an
    interleave over hash partitioned aggregates is re-derived as an interleave,
    and stays a union with prefer_existing_union = true.

The test harness's tree-node integrity block now applies the same pre-pass so
it mirrors the rule. Existing union_to_interleave and
union_not_to_interleave tests are unchanged, and union.slt passes.

Are there any user-facing changes?

No API changes. One behavioral note: with
datafusion.optimizer.prefer_existing_union = true, an InterleaveExec
present in the input plan is now emitted as a UnionExec. The rule never
creates interleaves under that setting, and plans produced by the default
pipeline are unaffected because EnsureRequirements is the only creator of
InterleaveExec and runs once.

@github-actions github-actions Bot added optimizer Optimizer rules core Core DataFusion crate labels Sep 6, 2026
@jayzhan211

Copy link
Copy Markdown
Contributor Author

@zhuqi-lucas WDYT about converting the interleave into a union like this? Are there concerns I'm missing?

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 88.88889% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 81.62%. Comparing base (3dfa245) to head (7f49582).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
.../physical-optimizer/src/ensure_requirements/mod.rs 0.00% 0 Missing and 1 partial ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main   #24959      +/-   ##
==========================================
- Coverage   81.62%   81.62%   -0.01%     
==========================================
  Files        1124     1124              
  Lines      412462   412474      +12     
  Branches   412462   412474      +12     
==========================================
- Hits       336684   336682       -2     
- Misses      55965    55973       +8     
- Partials    19813    19819       +6     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@zhuqi-lucas zhuqi-lucas left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @jayzhan211 — approach matches what #21827 landed on.

One scope note: the invariant itself is unchanged. On your branch, rebuilding an interleave whose children lost their partitioning (what any transform_up does) still fails:

Internal error: Assertion failed: can_interleave(children.iter())

EnsureRequirements is safe because Phase 0 strips interleaves before anything rebuilds them, but other rules aren't. JoinSelection runs at optimizer.rs:101, before this rule at :136, so the case reported in #21826 — a custom chain running EnforceDistribution twice, JoinSelection then walking an interleave the earlier pass built — isn't covered here. Same goes for the ~10 rules after :136 that rebuild the tree while an interleave is live.

Not blocking — the idempotency fix is worth having. Just maybe Part of #21826 rather than Closes.

@jayzhan211

Copy link
Copy Markdown
Contributor Author

Let me find another solution to close #21827

@jayzhan211
jayzhan211 added this pull request to the merge queue Sep 6, 2026
Merged via the queue into apache:main with commit c14633e Sep 6, 2026
42 checks passed
@jayzhan211
jayzhan211 deleted the enforce-dist branch September 6, 2026 06:08
lyne7-sc pushed a commit to lyne7-sc/datafusion that referenced this pull request Sep 7, 2026
…ecutable` (apache#24963)

## Which issue does this PR close?

- Closes apache#21826. Builds on apache#24959, which handled the
`EnsureRequirements` side.

## Rationale for this change

`InterleaveExec` is only valid while every child shares the same hash
(or
range) partitioning. `EnsureRequirements` creates it from a `UnionExec`
when
that happens to hold. Any later rule that rewrites a child and changes
its
output partitioning, for example `JoinSelection` swapping a join's
sides,
makes the optimizer's tree walk rebuild the interleave over children
that no
longer match. `InterleaveExec::replace_children` asserted on that, so
the
whole rule failed:

```
join_selection
caused by
Internal error: Assertion failed: can_interleave(children.iter()):
Can not create InterleaveExec: new children can not be interleaved.
```

This is the chain reported in apache#21826: distribution enforcement builds
the
interleave, `JoinSelection` breaks it, and no rule can repair it because
the
failure happens inside the rebuild, before any rule's closure runs.

apache#24959 made `EnsureRequirements` itself immune by normalizing
interleaves back
to unions before it runs, but that does not help when the rebuild
happens
inside another rule.

The root cause is that `InterleaveExec` enforces a cross-child
distribution
property one level earlier than the rest of the framework does.
Co-partitioning
for joins is validated by
`InputDistributionRequirements::check_invariants`
only at `InvariantLevel::Executable`; a hash join rebuilt over children
that
lost their partitioning does not error at rebuild time, it becomes
temporarily
unsatisfied and is either repaired by a later distribution pass or
rejected by
`SanityCheckPlan`. This PR makes `InterleaveExec` behave the same way.

This is not the fallback proposed in apache#21827. Nothing is substituted for
a
different node and no error is downgraded to a log. The node stays an
`InterleaveExec`, the error is the same, and it is raised at the
executable
checkpoint instead of at rebuild time. The decision on how to recover is
left
to the optimizer rule, which apache#24959 already implements.

## What changes are included in this PR?

In `datafusion/physical-plan/src/union.rs`:

- `InterleaveExec::try_new` still asserts interleavability, so explicit
  construction and proto decoding keep the strict check.
- A private `try_new_unchecked` builds the node without the check.
  `replace_children` in `Recompute` mode now uses it.
- `compute_properties` only claims the shared hash / range partitioning
when
every child still has it, and reports `UnknownPartitioning` otherwise,
so
  nothing downstream can rely on a layout the node does not have.
- `InterleaveExec::check_invariants` runs the `can_interleave` check at
  `InvariantLevel::Executable`, with the same message as before.

The physical planner already checks `Always` after each rule and
`Executable`
at the end, so an unrepaired interleave is still rejected, with the same
error, by `SanityCheckPlan`.

## What is the testing strategy for this PR?

- `test_interleave_rebuild_defers_partitioning_invariant` (unit test in
`union.rs`): rebuilding over one round-robin child succeeds, `try_new`
on
the same children still fails, the node reports unknown partitioning,
the
`Always` check passes, and the `Executable` check fails with the
original
  message.
- `issue_21826_join_selection_after_distribution_pass`
  (`core/tests/physical_optimizer/enforce_distribution.rs`): the exact
reported chain. A union of two partitioned hash joins with statistics
chosen
  so `JoinSelection` swaps only one of them. On `main` this fails inside
`JoinSelection` with the error above. Here `JoinSelection` completes,
the
  interleave reports unknown partitioning, `SanityCheckPlan` rejects the
unrepaired plan, and a second `EnsureRequirements` pass produces a valid
  union that the sanity checker accepts (snapshotted).
- `interleave_broken_by_later_rewrite_is_repaired_on_next_pass`: the
same
mechanism where the disruption is a removable `CoalescePartitionsExec`,
so
the repair restores the interleave rather than demoting it, showing the
  optimization is not permanently lost.

## Are there any user-facing changes?

No API changes. Behaviorally, an `InterleaveExec` whose children lose
their
shared partitioning during optimization is now reported as an invariant
failure by `SanityCheckPlan` at the end of the pipeline instead of as an
error
from the rule that rebuilt it. The diagnostic is therefore coarser (it
names
the node, not the rule), which is the trade-off for letting rules
complete and
a later distribution pass repair the plan. Custom pipelines that run
rules
between two distribution passes no longer need a workaround.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate optimizer Optimizer rules v56.0.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants