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

fix: normalize InterleaveExec to UnionExec before enforcing distribution#24959
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
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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants