andygrove opened a new pull request, #2062:
URL: https://github.com/apache/datafusion-ballista/pull/2062
# Which issue does this PR close?
Closes #2047.
# Rationale for this change
With adaptive execution enabled (`ballista.planner.adaptive.enabled=true`),
15 TPC-DS queries fail with:
```
Internal("Assertion failed: partition < self.partitions:
EmptyExec invalid partition 1 (expected less than 1)")
```
The adaptive planner's `PropagateEmptyExecRule` collapses a provably empty
sub-plan into an `EmptyExec` that inherits the partition count of the node it
replaced, and the scheduler sizes a stage's task count from that count.
`datafusion-proto` encodes an `EmptyExec` as its schema alone —
`EmptyExecNode` has no partitions field, and the decoder rebuilds the node with
`EmptyExec::new(schema)`, which defaults to one partition. So a stage built
from an `EmptyExec` with N partitions launches N tasks against a plan that,
once it reaches the executor, can only execute partition 0. Every task above
partition 0 then trips the assertion.
The static planner is unaffected because it never produces a multi-partition
`EmptyExec`: its only `EmptyExec` comes from `LogicalPlan::EmptyRelation`,
which has a single partition.
I filed the underlying serialization gap upstream as
[apache/datafusion#23642](https://github.com/apache/datafusion/issues/23642).
# What changes are included in this PR?
- `make_empty_exec_serde_safe` in `ballista/scheduler/src/planner.rs`
rewrites every multi-partition `EmptyExec` into a round-robin `RepartitionExec`
over a single-partition `EmptyExec`. The plan is equivalent, and a
`RepartitionExec` carries its partitioning through serialization, so the count
arrives on the executor intact.
It is applied in `create_shuffle_writer_with_config`, the single point
where a stage plan is built for the wire. That is after all physical optimizer
rules, so no later rule can collapse the rewrite, and it is a no-op for the
static planner.
- A unit test for the rewrite and for the single-partition case, asserted on
the rendered plan string.
- A test in `ballista-core` pinning the upstream `datafusion-proto`
behaviour that makes the workaround necessary. When that test starts failing,
the partition count survives serialization upstream and
`make_empty_exec_serde_safe` can be deleted.
- Two AQE stage-plan snapshots updated. Both cover a stage plan carrying a
2-partition `EmptyExec` — exactly the shape that fails on the wire today.
# How are these changes tested?
`cargo test --workspace` and `cargo clippy --all-targets --workspace -- -D
warnings` are clean.
Verified end-to-end against SF1 TPC-DS on a native 1-scheduler / 1-executor
cluster (8 slots), with `ballista.planner.adaptive.enabled=true` and
`datafusion.optimizer.prefer_hash_join=false`, running each of the 15 queries
listed in #2047. The assertion no longer occurs on any of them.
13 of the 15 (q8, q16, q24, q29, q34, q37, q41, q43, q46, q54, q61, q68,
q73) now run to completion. q1 and q44 no longer hit the assertion but do not
complete: they hang on a **separate, pre-existing** defect that the assertion
was previously masking. Once the residual plan collapses to empty, the replan
removes an in-flight stage's exchange and cancels the stage; the tasks already
running then complete, and the late status update fails `finalise_stage` with
`Execution("Can't find active cache resolve")`, so the job never finishes. That
error is visible in the scheduler log on `main` today, before this change. It
is unrelated to serialization and is left for a follow-up issue.
Also unchanged by this PR, and tracked separately: the empty results
themselves fall under the AQE correctness divergences in #2046.
--
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]
