feat: distributed EXPLAIN, EXPLAIN FORMAT TREE, and EXPLAIN ANALYZE (rebased onto spiceai-52.5)

[phillipleblanc]

Reimplements #31 on top of spiceai-52.5 (DataFusion 52). Supersedes #31.

What

End-to-end support for the three EXPLAIN forms over distributed Ballista:

• EXPLAIN ... — adds a distributed_plan row alongside logical_plan and physical_plan.
• EXPLAIN FORMAT TREE ... — physical plan rendered with box-drawing tree characters; logical_plan is omitted to match DataFusion's native behavior.
• EXPLAIN ANALYZE ... — a single Plan with Metrics row containing the per-stage, per-operator metrics gathered from the scheduler after the job completes.

How

EXPLAIN / FORMAT TREE — ExplainFormat doesn't survive a datafusion-proto round-trip, so the client wraps LogicalPlan::Explain in a new BallistaExplainNode logical extension before sending to the scheduler. The scheduler unwraps it back into a native LogicalPlan::Explain (unwrap_ballista_explain in state/mod.rs) and then its existing physical-planning intercept replaces ExplainExec with a distributed-aware variant. extract_logical_and_physical_plans and construct_distributed_explain_exec are threaded with &ExplainFormat so Tree/Indent/PgJson/Graphviz are honored.

EXPLAIN ANALYZE —

• Client (BallistaQueryPlanner): strips LogicalPlan::Analyze, runs the inner plan via DistributedQueryExec, and wraps it in a new DistributedExplainAnalyzeExec. DistributedQueryExec gains an Arc<Mutex<Option<String>>> job_id handle that is populated when the job is submitted; the parent reads it after the child stream drains and calls the new GetJobMetrics RPC.
• Scheduler: new GetJobMetrics RPC (proto: GetJobMetricsParams/Result, JobStageMetrics, OperatorWithMetrics, OperatorMetric) walks the execution graph in the same pre-order DFS order as ballista_core::utils::collect_plan_metrics so the per-operator metrics line up with the rendered plan text. Falls back from the active-job cache to the saved completed-job graph (get_job_execution_graph) so the call still succeeds after succeed_job moves the graph out of active_job_cache.

Tests

ballista/client/tests/context_checks.rs:

• should_execute_explain_query_correctly (was already there)
• should_execute_explain_format_tree_query_correctly — checks two-row result (no logical_plan) and box-drawing characters in physical_plan.
• should_execute_explain_analyze_query — sanitizes metric values to ... and checks both stages, both ShuffleWriterExec and AggregateExec, and that per-operator metrics=[...] appear.

All three cases run in both standalone and remote modes (the remote mode goes over the gRPC scheduler, exercising the codec round-trip and the GetJobMetrics RPC):

test supported::should_execute_explain_query_correctly::case_1_standalone ... ok
test supported::should_execute_explain_query_correctly::case_2_remote ... ok
test supported::should_execute_explain_format_tree_query_correctly::case_1_standalone ... ok
test supported::should_execute_explain_format_tree_query_correctly::case_2_remote ... ok
test supported::should_execute_explain_analyze_query::case_1_standalone ... ok
test supported::should_execute_explain_analyze_query::case_2_remote ... ok

Also manually validated against a local distributed cluster (scheduler + executor with mTLS) by issuing the queries through Spice's HTTP API:

• EXPLAIN returns the expected three rows.
• EXPLAIN FORMAT TREE returns two rows with tree-formatted physical plan.

Notes / follow-ups

• This PR covers the ballista client API (BallistaQueryPlanner) path. Spice's /v1/queries async path submits the raw LogicalPlan::Analyze directly via scheduler.submit_job, bypassing BallistaQueryPlanner, and the scheduler currently only intercepts Explain, not Analyze. A follow-up is needed in either Spice or in this scheduler submit_job path to get EXPLAIN ANALYZE working over /v1/queries.

[Copilot]

Pull request overview

Adds end-to-end distributed support for EXPLAIN, EXPLAIN FORMAT TREE, and EXPLAIN ANALYZE in Ballista (DataFusion 52), ensuring explain-format survives the client→scheduler proto round-trip and enabling scheduler-sourced per-stage/operator metrics for analyze output.

Changes:

• Introduces BallistaExplainNode logical extension + codec/proto support to preserve ExplainFormat across datafusion-proto, and unwraps it on the scheduler back into native LogicalPlan::Explain.
• Threads ExplainFormat through distributed explain rendering, matching DataFusion behavior for FORMAT TREE (omit logical_plan row; tree-rendered physical_plan).
• Adds GetJobMetrics scheduler gRPC RPC and a client-side DistributedExplainAnalyzeExec that fetches and renders per-stage/operator metrics as Plan with Metrics.

Reviewed changes

Copilot reviewed 13 out of 14 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
ballista/scheduler/src/state/task_manager.rs	Makes execution-graph lookup available for gRPC metrics retrieval.
ballista/scheduler/src/state/mod.rs	Unwraps BallistaExplainNode, captures ExplainFormat, and threads it into distributed explain exec construction.
ballista/scheduler/src/state/distributed_explain.rs	Formats explain output according to ExplainFormat and matches DataFusion’s FORMAT TREE row behavior.
ballista/scheduler/src/scheduler_server/grpc.rs	Adds GetJobMetrics RPC and stage/operator metrics serialization using plan traversal.
ballista/executor/src/execution_loop.rs	Updates test SchedulerGrpc mock to include the new RPC.
ballista/core/src/serde/mod.rs	Extends Ballista logical extension codec to encode/decode BallistaExplainNode via Ballista protobuf.
ballista/core/src/serde/generated/ballista.rs	Regenerates protobuf bindings for ExplainNode and GetJobMetrics messages/RPC.
ballista/core/src/planner.rs	Wraps EXPLAIN for distribution and implements distributed EXPLAIN ANALYZE via DistributedExplainAnalyzeExec.
ballista/core/src/extension.rs	Defines BallistaExplainNode and stable string mapping for explain formats.
ballista/core/src/execution_plans/mod.rs	Registers the new distributed_explain_analyze execution plan module.
ballista/core/src/execution_plans/distributed_query.rs	Adds a shared job_id handle so parents can query scheduler metrics post-execution.
ballista/core/src/execution_plans/distributed_explain_analyze.rs	New exec that runs a distributed job, then fetches and renders metrics as a single-row EXPLAIN ANALYZE result.
ballista/core/proto/ballista.proto	Adds LogicalPlanExplainNode and GetJobMetrics RPC + message types.
ballista/client/tests/context_checks.rs	Adds integration coverage for EXPLAIN FORMAT TREE and EXPLAIN ANALYZE in standalone + remote modes.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 18 out of 19 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 21 out of 23 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.