jayzhan211 opened a new pull request, #24959:
URL: https://github.com/apache/datafusion/pull/24959
## Which issue does this PR close?
- Closes #21826.
## 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
`RepartitionExec`s 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:
```rust
// 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
