apache / apache/datafusion-ballista
Allow a stage to start on partially complete inputs instead of waiting for the whole upstream stage
- Dominant language
- Rust
- Stars
- 2.1k
- Forks
- 320
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 66
Description
**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**
A downstream stage cannot start until every one of its input stages is complete:
```rust
// ballista/scheduler/src/state/execution_stage.rs
pub fn resolvable(&self) -> bool {
self.inputs.iter().all(|(_, input)| input.is_complete())
}
```
So a stage finishes when its slowest task finishes, and until then no downstream work runs even if 99% of the input has been written and the cluster has idle slots. On a skewed key or a slow node this is the single largest source of avoidable latency in Ballista today.
The important word is *avoidable*. Waiting on the whole upstream stage is a property of the current scheduler, not of materializing stage output. The shuffle files for the finished producer tasks already exist and are already addressable; nothing about the blocking model requires a consumer to wait for the rest of them.
**Describe the solution you'd like**
Let a consumer stage start on partially complete inputs when there is spare capacity, learning about additional `PartitionLocation`s as the remaining producer tasks finish.
Rough shape:
- Relax the resolvability condition from "all inputs complete" to a policy decision: enough input available, and free slots that would otherwise sit idle.
- `ShuffleReaderExec` (and `RangeShuffleReaderExec`) gain a way to take on locations after the plan is resolved, or the consumer task is launched over the known subset with the remainder scheduled as follow-up work.
- Keep files, retries, and slot accounting exactly as they are, so `FetchPartitionError` handling, `rollback_running_stage`, and `rerun_successful_stage` continue to work.
- Guard against starving the producer: consumers must not take slots the remaining producer tasks need, or the query gets slower rather than faster.
Constraints that need care:
- AQE. Rules that need a completed stage (notably empty-stage elimination, and coalescing based on exact per-partition counts) either have to run before the early start or be skipped for stages that start early. Deciding which of the two happens is most of the design work.
- Per-task plan rewriting, which already has to satisfy several competing constraints around partition counts.
**Describe alternatives you've considered**
- Bubble execution (#408) and all-at-once scheduling (#1151) are the general form of this: decide how much of the DAG is in flight at once rather than always running exactly one stage. This issue is the narrow, incremental version that keeps one-stage-at-a-time scheduling as the default and only overlaps the tail.
- Fully pipelined exchange (#2003). Bigger change, gives up the properties documented in the shuffle design page.
- Attack skew directly instead (better partitioning, splitting oversized partitions). Complementary: it shrinks the stall rather than overlapping it.
**Additional context**
- Design rationale and the list of properties any change here should preserve: `docs/source/contributors-guide/shuffle.md` (#2308).
- Raised by @Dandandan on #2308.
Contributor guide
Research direction
Start with ballista/scheduler/src/state/execution_stage.rs and the resolvable condition, then read the ShuffleReaderExec and RangeShuffleReaderExec areas mentioned in the issue. Review docs/source/contributors-guide/shuffle.md and the AQE constraints before choosing how partially complete inputs are represented. Done means downstream work can overlap the upstream tail without breaking retries, slot accounting, or required AQE behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100