apache / apache/datafusion-ballista

Emit the final (post-AQE) physical plan to the event log

Open
#2,422 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
2.1k
Forks
320
Avg merge
1d 22h
Merged PRs (30d)
66

Description

_Drafted with an LLM (Claude Code) from a review comment on #2419; reviewed by me before posting._

Follow-up to https://github.com/apache/datafusion-ballista/pull/2419#issuecomment-5532791897 — @milenkovicm: *"maybe we could emit final plan to event log when AQE used"*.

### What's there today

Verified against `d9fadd8c`:

- **`JobStart.physical_plan`** — `graph.physical_plan()` rendered at submission (`scheduler_server/event_log.rs:53`). For AQE this is the plan after `AdaptivePlanner::try_new`'s first optimizer pass, before any stage ran.
- **`JobEnd.stages[].stage_plan`** — per stage, frozen at dispatch: `create_resolved_stage` (`state/aqe/mod.rs:208`) stores the `ShuffleWriter` that `BallistaAdapter` produced at the moment the stage became runnable. This is the ground truth for what each stage actually executed.
- **`JobEnd.job.physical_plan`** — `graph_to_job_response` also renders `graph.physical_plan()` (`api/dto_build.rs:95`). For `AdaptiveExecutionGraph` that getter returns the live `planner.plan` (`state/aqe/mod.rs:714`), so this is *already* the post-AQE plan; for `StaticExecutionGraph` it returns the frozen submitted plan (`state/execution_graph.rs:690`).

So most of the bytes exist. Three things are wrong with relying on that:

1. **Same field name, two meanings, no marker.** Nothing in a log says "this job used AQE", or that the two `physical_plan` strings are a before/after pair for adaptive jobs and the same plan twice for static ones. The plan-dump tooling behind #2419 read only `JobStart.physical_plan` and never surfaced the final one.
2. **Only the terminal state is recorded.** Every intermediate replan is lost — join flips from `SelectJoinRule` / `DemoteUnsafeBroadcastJoinRule`, branch pruning from `PropagateEmptyExecRule`, `CoalescePartitionsRule` decisions. Answering "why did this stage become a `CollectLeft`" needs the revision that made the call, not the end state.
3. **The final plan is not necessarily a faithful record of what ran.** `replan_stages` (`state/aqe/planner.rs:356`) re-runs the whole optimizer chain over the entire tree after every stage completion, and `ExchangeExec::with_new_children` (`state/aqe/execution_plan/exchange.rs:426`) carries the resolved shuffle state across a rewrite — so a subtree under an already-resolved exchange can be rewritten *after* that stage executed. When that happens the whole-query "final plan" and the frozen `stage_plan` disagree, and only the latter is right.

### Proposed design

The constraint that shapes this: the history server reads **only** `JobEnd` (`ballista/history/src/reader.rs:150`) and relays `JobEnd.job` / `JobEnd.stages` verbatim. A new *sibling* field on `JobEnd` is invisible to `GET /api/job/{id}` unless both the reader and the serving path change. A field added *inside* `JobResponse` flows to the live API and the replayed API for free and preserves the byte-identical-replay property.

**1. Make the two plans explicit on `JobResponse` (`ballista-api-types`)**

- `physical_plan` — plan as submitted, same meaning for both graph kinds
- `final_physical_plan: Option` — post-AQE plan; `None` for non-adaptive jobs
- `adaptive: bool`, `replans: u32`

Populated in `graph_to_job_response`.

**2. Split the getter on the `ExecutionGraph` trait**, so the two implementations stop meaning different things by one name:

- `physical_plan()` keeps returning the *current* plan
- add `submitted_physical_plan(&self) -> Option<&str>` — rendered and frozen at construction, symmetric with the existing `logical_plan()`. `AdaptiveExecutionGraph::try_new` renders it right after `AdaptivePlanner::try_new`; `StaticExecutionGraph` returns its one plan.
- add `fn is_adaptive(&self) -> bool { false }` and `fn replan_count(&self) -> u32 { 0 }` with defaults; the adaptive graph overrides both, bumping the counter in `replan_stages`.

**3. No event-log schema change.** No `SCHEMA_VERSION` bump: the new fields ride inside the existing `JobEnd.job` payload, which is stored as opaque raw JSON precisely so this class of change is free. `JobStart.physical_plan` keeps its current meaning.

**4. (opt-in, separate PR) per-replan history.** New event kind `JobPlanRevision { revision, trigger_stage_id, at, plan }`, appended when a replan changes the plan, behind a scheduler flag (`--event-log-plan-revisions`, default off). Unknown record kinds are already skipped by readers and `testdata/schema-v1.eventlog` freezes that behaviour, so this lands without touching compatibility. Default off because it is expensive — q8 at SF1000 has 8 stages and `DataSourceExec` inlines the full parquet file list per file group, so a full render per replan multiplies an already-large payload (also raised in #2419).

### Alternatives considered

- **Sibling field on `JobEnd`.** Simplest to write, but invisible through the history server's REST surface without also changing the reader and the serving path, and it breaks the "relay stored bytes verbatim" property. Rejected.
- **Redefine `JobEnd.job.physical_plan` as "final" and just document it.** Zero code, but leaves adaptive and static meaning the same field differently, and gives no AQE marker. Rejected.
- **Structured plan (JSON/protobuf) instead of rendered text.** Better for tooling, much bigger change, inconsistent with every other plan field in the log. Out of scope.

### Open questions

1. Should `JobResponse.physical_plan` be redefined to "as submitted" for adaptive jobs — a behaviour change to the live API, which today returns the current plan — or should the new field be `submitted_physical_plan` and `physical_plan` left alone?
2. Is a rewrite below an already-resolved `ExchangeExec` actually reachable, or is it prevented somewhere I haven't found? If reachable, should the final-plan render mark subtrees that no longer match the `stage_plan` that ran?
3. Worth capturing the final plan on the failure and cancel paths too (`job_cancel_event`, `JobRunningFailed`)? The plan at the point of failure seems more useful than the submitted one, but the graph is mid-flight there.
4. Is the per-replan history (4) wanted at all, or are the final plan plus the per-stage `stage_plan`s enough?

Contributor guide

Open the contributing guide

Research direction

Start with JobResponse and graph_to_job_response in api/dto_build.rs, then inspect the ExecutionGraph implementations in state/aqe/mod.rs and state/execution_graph.rs. Read ballista/history/src/reader.rs and testdata/schema-v1.eventlog before resolving the open API and AQE semantics questions. Done means the submitted and final plan metadata is consistently exposed through live and replayed job responses, with appropriate tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
api, distributed-systems
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.