NoahKusaba opened a new pull request, #2459:
URL: https://github.com/apache/datafusion-ballista/pull/2459
## Rationale for this change
AQE decides to broadcast a join's build side from *estimated* statistics.
When that side then *measures* larger than the probe side, the replan's
`join_selection` swaps the `CollectLeft` join and the broadcast `ExchangeExec`
lands on the probe side. A broadcast reader reports a single partition, and a
`CollectLeft` join has as many partitions as its probe side, so the whole join
stage collapses to one task.
TPC-H q8 at SF1000 hit this in stage 5: one task, 79.9M input rows, `min =
median = max = 7429ms`, feeding a 256-way shuffle write.
```
SortShuffleWriterExec: partitioning=Hash([l_orderkey@0], 256)
HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(s_suppkey@0,
l_suppkey@1)], ...
CoalescePartitionsExec
ShuffleReaderExec: upstream_stage: 3, partitioning:
UnknownPartitioning(256)
ShuffleReaderExec: upstream_stage: 2, broadcast: true,
upstream_partition_count: 256
```
The probe-side reader should be `partitioning: UnknownPartitioning(256)`, so
the stage spreads across executors instead of one core.
## What changes are included in this PR?
- **`ExchangeExec::as_probe_side`** implements the DataFusion hook added in
<DATAFUSION PR LINK>. When a `CollectLeft` swap moves a broadcast exchange to
the probe side, it returns the same stage read partitioned. Broadcasts over
ordered inputs return `None` and are left to the k-way merge reader.
- **`ExchangeExec::to_partitioned`**, the inverse of the existing
`to_broadcast`: the same exchange with `broadcast` off, sharing its stage state
and keeping its plan id so it replaces the original in place.
- **`take_stage_output_partitions`** stores a broadcast stage's locations
one entry per partition, in partition order, sized by the stage's *input*
partition count. The old path asked the exchange, which reports `1` for a
broadcast, and stored the locations unordered from a `HashMap`. That was fine
while the only reader merged everything into one partition; a partitioned read
needs entry `i` to be partition `i`.
- **`StageOutput::partition_locations_broadcast`** is removed.
`partition_locations(n)` now serves both readers, and this was its only caller.
## What is the testing strategy for this PR?
Three tests in `scheduler/src/state/aqe/test/probe_side_broadcast.rs`, each
covering a distinct failure. Reverting the partition-count fix alone fails only
the second, so they do not overlap:
- `swapped_broadcast_is_read_partitioned_on_the_probe_side` — the swapped
join's probe side reads 4 partitions, not 1.
- `swapped_broadcast_reads_the_locations_its_tasks_reported` — a completed
broadcast stage stores one entry per partition, keyed by partition id rather
than arrival order.
- `swapped_null_aware_broadcast_is_coalesced_on_the_probe_side` — a
null-aware `NOT IN` join still swaps to `RightAnti` and keeps its single-task
probe side.
## Are there any user-facing changes?
No API changes. Join stages whose broadcast side is swapped onto the probe
side now run across all partitions instead of a single task.
## Note
This depends on `ExecutionPlan::as_probe_side` landing in DataFusion
https://github.com/apache/datafusion/pull/25354 and on Ballista upgrading to a
release that contains it. Until then the branch does not build against the
released DataFusion.
--
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]
