Support `EXPLAIN ANALYZE` in Ballista

[danielhumanmod]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

The initial EXPLAIN support was merged in #1293 , but the EXPLAIN ANALYZE portion is not implemented yet. This issue tracks the remaining work needed to support it.

Describe the solution you'd like
The exact design is still TBD, but the general direction is to create a new exec to gather per-executor metrics from scheduler so we can generate a full EXPLAIN ANALYZE output

Describe alternatives you've considered
A clear and concise description of any alternative solutions or features you've considered.

Additional context
Add any other context or screenshots about the feature request here.

[milenkovicm]

Would you be interested to investigate this @danielhumanmod? I got stacked with some other functionality I'd like to add to ballista ?

[danielhumanmod]

Yes I will take on look on this, thanks for looping me in

[danielhumanmod]

An alternative would be to add ability to register job completion callback which triggers change of job B once the job A finishes. Trigger can change UnresolvedShuffle to actual Exec (we could use same logic like we use for simple EXPLAIN).
I believe job finish callback would be useful for other cases, I'm not sure at this point how complicated would it be to implement, may have a look this weekend. (#1309 (comment))

Hey @milenkovicm I did some investigation on this approach. The current processing logic for QueryStageSchedulerEvent::TaskUpdating even is already behaves like a built-in callback chain:

TaskUpdating event
  → update_task_statuses
  → update_task_status
     → update_stage_output_links
     → processing_stages_update
         → resolve_stage
             → remove_unresolved_shuffles

This whole chain is tightly coupled, extracting it into a standalone job-completion callback looks non-trivial

But I find remove_unresolved_shuffles() is actually the point where UnresolvedShuffleExec -> ShuffleReaderExec, so another reasonable approach might be adding a lightweight abstraction for controlling how shuffles get resolved. Something like a ShuffleResolver that lets us plug in different UnresolvedShuffleExec replace strategies:

trait ShuffleResolver {
fn should_resolve(...): bool;
fn resolve(...): Result<Arc<dyn ExecutionPlan>>;
}

The flow will looks like:

remove_unresolved_shuffles()
  ↓
resolver_manager.resolve()
  ↓
select_resolver()  
  ├─ AnalyzeShuffleResolver.should_resolve()?
  │  ├─ true → Use this resolver and stop
  │  └─ false → continue
  │
  ├─ DefaultShuffleResolver.should_resolve()?
  │  └─ true → use this and stop
  │
  ...

[milenkovicm]

thank you @danielhumanmod for tackling this issue, I think it would bring few new possibilities if implemented .

I'm not sure if we should utilise shuffles as base building block. Let's say we want to implement distributed cache functionality or some kind of job preparation, they may be reading from source which is not a shuffle file.

Would it be possible to keep job submission pending until another job finishes? If we can attach a callback on a "leading" job which would change dependant job status from "pending" to "ready to run" or to "cancel".

I'm just throwing ideas, will have to take a deeper look at this

[danielhumanmod]

Thanks for the insight, @milenkovicm — job-level dependency is a really good idea and definitely aligns well with where the scheduler could evolve. It also fits nicely on top of the current stage/shuffle-level dependency logic.

I took a quick look at the implementation, and the main changes I’m thinking about are:
1. If needed, split jobs and build the dependency context in the scheduler (execute_query()).
2. Add job-level dependency tracking in TaskManager for pending jobs.
3. Trigger downstream tasks on JobFinished events.

I’ll start with a small PoC around this, but happy to adjust the direction based on your feedback. Sorry for the slow response — I was AFK for a bit over the holidays. Hope you had a nice break as well :)