andygrove opened a new issue, #2652: URL: https://github.com/apache/datafusion-ballista/issues/2652
> Disclaimer: this issue was drafted by an LLM (Claude Code) at my request, against `b89a4e46`. File/line references were checked against the code, but please treat the design as a proposal and push back where it's wrong. ## Summary This proposes an opt-in execution mode, `ballista.execution.mode = pipelined`. In this mode the scheduler places every task of a job up front, and stages stream record batches to each other over Arrow Flight instead of running stage by stage through shuffle files. This is the model used by Trino (with `retry_policy=NONE`) and by DataFusion Distributed. The goal is lower latency for interactive queries. The blocking shuffle stays the default and is unchanged. Background: - the discussion in #2003; - #1151, which proposed streaming exchange with all-at-once scheduling; - bubble execution (#408), the per-edge generalization; - the costs of today's model, set out in the [shuffle design page](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/docs/source/contributors-guide/shuffle.md). ## Why support both models Ballista should support both models for different workloads, rather than replace one with the other. - **Batch and ETL need the blocking model.** The properties in the shuffle design page matter most for large batch and ETL jobs: - queries wider than the cluster; - re-running only what was lost after a failure; - executors joining and leaving between stages; - exact statistics for AQE; - intermediates larger than memory. - **Interactive queries are better served by the pipelined model.** Where those properties don't matter, it wins: short queries with small intermediates, where per-stage materialization and the stage barrier dominate, and where time to first row matters. - **Most of the system is shared.** SQL and physical planning, the stage split, executors, operators, serde, the REST API and the UI stay common. The new code is one scheduler graph, one exchange operator pair and a client read path. - **There is precedent.** Trino ships both models behind `retry_policy`. Its fault-tolerant mode, which writes exchange data to a spooling store, is essentially Ballista's model. Its default pipelined mode is the one proposed here. DataFusion Distributed's own documentation recommends Ballista for large batch and ETL work that benefits from materialized intermediates, and targets interactive analytics itself. - **Supporting both opens the way to mixing them.** Later, one query could use pipelined regions joined by blocking edges where those are needed (#408), which neither pure model can do. The cost is a second scheduling path to maintain and test. To contain it: - keep the mode opt-in and marked experimental; - share code paths wherever possible; - run the TPC-H correctness suite in both modes in CI. ## Comparison | | Blocking (today) | Pipelined (proposed) | | --- | --- | --- | | Cost at each stage boundary | Every intermediate byte is written, read and often fetched over the network; the next stage waits for the slowest task | Batches stream in memory; downstream stages run alongside upstream ones | | Time to first row | After the whole job completes | As soon as the final stage produces output | | Intermediates larger than memory | Bounded by local disk | Bounded by memory with backpressure; no spill at first | | Query wider than the cluster | Works; tasks queue for vcores | Must fit; waits for capacity, and is rejected if it can never fit | | Executor or task failure | Re-runs only what was lost | The query fails (whole-query retry could come later) | | Executors joining or leaving | Allowed between stages | Fixed for the life of the query | | Runtime re-optimization (AQE) | Exact statistics at each stage boundary | Static plan (sampling-based adaptivity could come later) | | Sharing the cluster with other jobs | Fine-grained, task by task | Capacity reserved for the life of the query | | Local storage | Required for shuffle | Not used by the exchange | | Dynamic filters across stages | Disabled today (#1375) | A natural follow-up | ## Proposed architecture ### Configuration - `ballista.execution.mode` is a session setting: `blocking` (the default) or `pipelined`. - Pipelined jobs require the push-staged scheduler. - Under pull-staged, each executor claims work on its own poll, so a job can't be placed atomically. - A pipelined job submitted to a pull-staged scheduler is rejected at submission with a clear error. ### Planning - A new `PipelinedExecutionGraph` implements the existing [`ExecutionGraph`](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/scheduler/src/state/execution_graph.rs#L107) trait. It is chosen in [`submit_plan`](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/scheduler/src/state/task_manager.rs#L277), alongside the static and adaptive graphs. - It uses `DefaultDistributedPlanner`, so the stage cut points don't change: hash repartitions, coalesces, sort-preserving merges, broadcast build sides and union branches. - AQE is not used, because it plans a stage only once the stages before it have finished. - Stage roots become `PipelinedShuffleWriterExec`. - The re-optimization in [`to_resolved`](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/scheduler/src/state/execution_stage.rs#L513) is skipped. That step re-runs `JoinSelection` and `AggregateStatistics` against actual stage statistics, and pipelined stages have none yet. ### Placement - **All or nothing.** A partially placed pipelined job would deadlock. - When the push-mode bind pass offers resources, before the distribution policy runs, either every task of every stage is bound or none is and the job waits. - A job that could never fit in the cluster fails immediately. - **Task sizing follows today's rules.** A stage's partitions are split into tasks across executors, each task taking up to the executor's free vcores and at most `max_partitions_per_task`. A collapse stage gets a single task. - **Accounting.** On each executor, a job reserves the largest number of vcores that any one of its stages uses there, not the sum. - The job's stages share CPU, and a blocked task is a parked future rather than a thread. - The reservation is held until the job's tasks on that executor have finished. - **Wiring consumers to producers.** Once producers are placed, each consumer stage's `UnresolvedShuffleExec` is replaced with a `PipelinedShuffleReaderExec`. The reader lists, for each partition, the producer tasks and executors to read from. This mirrors how [`remove_unresolved_shuffles`](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/scheduler/src/planner.rs#L734) fills in file locations today. All tasks are then launched together. ### Exchange **Writer (`PipelinedShuffleWriterExec`).** It drives all of its task's input partitions concurrently and routes batches into per-output-partition queues. The queues live in an executor-wide registry, keyed by job, stage, task and partition. - Hash outputs are split with DataFusion's `BatchPartitioner` and coalesced back up to the session batch size before sending. Passthrough outputs keep the input partitioning. - Queues are bounded, so a slow consumer stalls its producer instead of growing memory. - Each partition ends with an explicit end-of-stream marker. If a producer fails or is cancelled, the consumer sees an error, never a silently truncated result. - Broadcast outputs, which are the build sides of `CollectLeft` joins, are retained so that every consumer task can replay them. They are released at job cleanup. Their size is already bounded by `ballista.optimizer.broadcast_join_threshold_bytes`. **Reader (`PipelinedShuffleReaderExec`).** It opens its streams when the consumer plan is executed, and drains all producers concurrently. - It can't reuse `ShuffleReaderExec`'s fetch path, which caps in-flight requests and buffers each remote partition whole so it can retry it ([shuffle_reader.rs#L781](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/core/src/execution_plans/shuffle_reader.rs#L781), [#L994](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/core/src/execution_plans/shuffle_reader.rs#L994)). - Capped reads deadlock against bounded producer queues. Consumer A waits on producer 1. Producer 1 is blocked on a full queue for consumer B. Consumer B is waiting on producer 2. - There are no retries: a broken stream fails the query. **Transport.** A new Flight `do_get` ticket action, `FetchPipelinedPartition`, is served from the registry by each executor's Flight service. - A request that arrives before the producer task has registered waits for a bounded time. - At first there is one stream per pair of producer task and consumer partition. - A follow-up could carry all of a consumer task's partitions over one stream per producer task, as DataFusion Distributed does. ### Start semantics - Tasks start as soon as they are launched. Backpressure, not scheduling order, keeps a probe-side producer from running far ahead of a join that is still building its hash table. - There are two alternatives, either of which could be adopted later if starting everything at once wastes CPU on the critical path: - DataFusion Distributed starts a producer only when its consumer first polls it. - Trino's default phased policy schedules build sides before probe sides. ### Client - Today `DistributedQueryExec` waits for `SuccessfulJob` and then fetches result partitions one after another ([distributed_query.rs#L812-L866](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/core/src/execution_plans/distributed_query.rs#L812-L866)). - With bounded queues that would deadlock, because the final stage can't finish until someone drains it. - For pipelined jobs: - Once the job is placed, the scheduler adds the final stage's locations to the `Running` status, which reaches the client over the existing `ExecuteQueryPush` stream. - The client then reads all partitions concurrently. - It also keeps watching the status stream, so a job failure surfaces even if no data stream breaks. ### Failure, cancellation and cleanup - **No stage or task retry.** If any task fails, or an executor running one of the job's tasks is lost, the job fails. Its remaining tasks are cancelled through the existing `CancelTasks` path. None of the `FetchPartitionError` rollback machinery applies. Re-running the whole query, like Trino's `retry_policy=QUERY`, is a possible follow-up. - **Early finish.** When the final stage completes, any tasks still running are cancelled. This is how `LIMIT` stops producers whose output is no longer needed. - **Cleanup.** Registry entries are dropped as each partition is fully consumed, and in any case at job cleanup through the existing `RemoveJobData` path. ### Memory - Exchange buffers are bounded by queue capacity, and per-task memory pools keep today's sizing. - There is no spill in the first version. A query whose operators exceed memory fails, just as it would in a single DataFusion process. ### Dynamic filters - These stay disabled in the first version, as they are today ([extension.rs#L757](https://github.com/apache/datafusion-ballista/blob/b89a4e4654b0bd5259856600d0e45703326b9d3f/ballista/core/src/extension.rs#L757), #1375). - The pipelined mode is where cross-stage filters become natural. A join's build side finishes while the probe-side producer is still waiting on backpressure, so a filter can still reach the producer's scan in time to matter. ## Scope of the first PR In scope: - everything above; - end-to-end tests in standalone mode: - TPC-H correctness in pipelined mode; - joins, broadcast joins, unions, sort-preserving merges and `LIMIT`; - cancellation and task failure; - docs. Follow-ups, out of scope for the first PR: - Per-edge exchange modes and pipelined regions (#408). These would let barrier-dependent features (AQE rules, range repartition cuts, `PrefixMergeExec`) keep a blocking edge inside a pipelined query. - Exchange buffers that spill to disk, along the lines of Flink's hybrid shuffle. - Cross-stage dynamic filters (#1375). - Multiplexed streams per task pair, and local reads when producer and consumer share an executor. - Whole-query retry, fair admission between pipelined and blocking jobs, and pull-staged support. ## Open questions 1. Is `ballista.execution.mode = blocking | pipelined` the right name and set of values? 2. Is reserving the largest stage's vcores on each executor the right accounting, or should a pipelined job reserve whole executors? 3. A pipelined job that doesn't fit right now could wait, at the risk of starving behind blocking jobs, or fall back to blocking execution. Which should it do? 4. Should producers start eagerly, or only when a consumer first polls them? -- 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]
