andygrove opened a new pull request, #2186:
URL: https://github.com/apache/datafusion-ballista/pull/2186
# Which issue does this PR close?
Closes #2184.
# Rationale for this change
Per-task restriction gives each branch of a `UnionExec` the partitions that
task owns and an empty slice to the rest, relying on an unowned branch becoming
0-partition so the union's index arithmetic routes past it.
That fails when a branch's partition count does not come from restrictable
leaves. A `CoalescePartitionsExec` or a non-preserving `SortExec` reports one
partition whatever its leaves do, and a broadcast `ShuffleReaderExec` is
deliberately never pruned. The branch keeps reporting a partition and keeps
producing all of its data, so the executor — which runs every partition of the
plan it is handed — runs every branch for every task.
At SF1 with `target_partitions=1`:
| Query | Expected | Got | |
|---|---|---|---|
| q5 | `73724.99` | `294899.96` | 4x |
| q80 | `11343.60` | `22687.20` | 2x |
| q49 | — | `Internal("FileStreamBuilder invalid partition index: 0")` | |
This does not show at higher partition counts because those union stages are
capped by a hash repartition, where the writer's output is an intrinsic K-space
and inner partition counts do not matter. Confirmed by instrumenting the
rewriter: restricting q49's 3-partition union stage to `[0]` returned a plan
that still had 3 partitions.
# What changes are included in this PR?
**`planner.rs` — cut a stage boundary beneath a union branch that cannot be
restricted away.** The branch becomes a `ShuffleReaderExec`, which restricts to
nothing cleanly, and the collapse moves into its own stage where it still sees
its whole input. Branches that already restrict away stay inline, so a union of
plain scans is unaffected and pays no extra shuffle.
Whether a branch can be restricted away is decided by *asking the rewriter*
(`restrict_plan_to_partitions(branch, &[])`) rather than by matching operator
types, so it stays correct as the rewriter evolves.
Branches holding a broadcast input or a `CollectLeft` join are treated as
unrestrictable regardless. Those build and probe sides can still be swapped
after stage planning, and the two orders restrict differently — the build side
is read whole, a broadcast input is never pruned — so a branch that looks
restrictable at planning time may not be by the time the task is built. That is
exactly what left q5's catalog branch running in every task after the first
version of this fix.
**`task_builder.rs` — take restriction scope from
`required_input_distribution()` for every operator** instead of a hardcoded
list of join types. `CrossJoinExec` collects its left input but was not on that
list, so each task of a multi-partition stage rebuilt the cross join's left
side from its own slice. This surfaced as an inflated q77 (whose catalog branch
is `from cs, cr`) once that branch got a stage of its own — caught by the
full-suite run, not by the original repro.
# Are there any user-facing changes?
No API change. Union queries whose branches collapse internally now
materialize those branches as their own stage, which costs one extra shuffle
per affected branch and is what makes the result correct.
## Testing
Four unit tests: a plain branch needs no stage of its own, a collapsing
branch and a broadcast-holding branch do, and a cross join's collected side is
read whole.
End-to-end at SF1 against a single-process DataFusion oracle:
- `--partitions 16` — **the full suite passes**, so the currently-green
configuration is unaffected.
- `--partitions 1` — q5, q49, q80, q61 and q77 all verify OK. Only q72
remains, and it exhausts its per-task memory share rather than returning a
wrong answer.
## Not covered
The adaptive planner builds stages through its own path (`state/aqe/`) and
is untouched here; q5, q80 and q61 still fail under
`ballista.planner.adaptive.enabled=true` at one partition. #2185 stays open,
and #2182 will still show the AQE legs red.
--
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]
