NoahKusaba opened a new pull request, #2818: URL: https://github.com/apache/datafusion-ballista/pull/2818
# Which issue does this PR close? Closes #2622. # Rationale for this change Table formats such as Iceberg and Delta finish a write with one commit, after every writer task has succeeded. Ballista sends every plan node to an executor, so today the commit runs as a one-task final stage on some executor. Executors then need the catalog's configuration and credentials, the commit must be idempotent because tasks are retried, and every executor caches catalog clients. Spark avoids this by committing on the driver. This PR gives Ballista the same shape: a job's final step can run on the scheduler, once, after all its stages have succeeded, and its small result becomes the job's result. # What changes are included in this PR? **Picking the root.** `SessionConfigExt::with_ballista_scheduler_root(predicate)` marks which root operators run on the scheduler. When a job's root matches, `TaskManager::plan_job` splits it off before the plan is cut into stages. This is the one split point for the static and adaptive planners, from logical or physical plans. A `CoalescePartitionsExec` with no row limit directly under the root is dropped, since the root reads all of its input's partitions anyway. **Running it.** `SchedulerRoot` (`state::scheduler_root`) holds the root. The execution graph (static and AQE) owns it as a `PendingRoot`. When the graph's last stage succeeds, it emits `RunSchedulerRoot` once instead of `JobFinished`, and the job stays Running. The event loop runs the root on a tokio task: - The root reads the final stage's output through a `ShuffleReaderExec` that always reads remotely. - If the job is no longer Running (for example, it was cancelled), the root is not run. - On success, the result is stored on the graph and the job succeeds. If the job was cancelled while the root ran, the result is discarded and the job stays cancelled. - If the root fails, the job fails with `Error running the job's scheduler root: ...`. **Returning the result.** The root's output is sent to the client in a new field, `SuccessfulJob.inline_result`, as Arrow IPC bytes. `DistributedQueryExec` decodes it in both push and pull mode. The Flight SQL service serves it as a local result through the new `QueryResult.batches`. **Limits.** A job with a root is refused: - if it runs pipelined, since the pipelined graph has no point where all stages have succeeded; - if it builds a `DataFrame::cache()`, since the cache would hold the root's input rather than the job's result. The result goes back in the job status, so it is capped at `ballista.client.grpc_max_message_size`. A commit's result is a few hundred bytes however much was written (a count is 392 bytes of IPC), so the cap only guards against a predicate that picks a root returning rows. If an executor holding the root's input is lost while the root runs, the job fails: the input is not rerun. **Cleanup.** For a job with an inline result, the final stage's shuffle data is freed with the other stages, since no client will fetch it. `clean_up_intermediate_job_data` is renamed `clean_up_spent_stage_data`. **UI and event log.** The job's displayed plan has the root on top. The REST stage list ends with a stage `"scheduler"` for the root, with no tasks or metrics. **Tests.** - Scheduler: the split, the refusals, the size cap, the stage list, cleanup, and the root's lifecycle (the job reports Running until the root has run; a job cancelled before the root starts never runs it; a job cancelled while it runs stays cancelled). - Client end-to-end (`ballista/client/tests/scheduler_root.rs`): a commit runs once on the scheduler under the static and adaptive planners, in push and pull mode; a rejected commit fails the job; a pipelined job is refused; an INSERT runs its sink on the scheduler. - Flight SQL end-to-end: a result computed on the scheduler is served to a Flight SQL client. # Are there any user-facing changes? - New config: `SessionConfigExt::with_ballista_scheduler_root`. Nothing changes for jobs that don't set it. - Proto: new field `SuccessfulJob.inline_result` (`proto-change`). - API: `ballista_flight_sql::QueryResult` has a new public field, `batches`, which breaks code that builds it with a struct literal (`api-change`). - REST: the stage list of a job with a scheduler root ends with a stage whose id is `"scheduler"` rather than a number. -- 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]
