Struct ExecutionGraph 

pub struct ExecutionGraph { /* private fields */ }

Represents the DAG for a distributed query plan.

A distributed query plan consists of a set of stages which must be executed sequentially.

Each stage consists of a set of partitions which can be executed in parallel, where each partition represents a Task, which is the basic unit of scheduling in kapot.

As an example, consider a SQL query which performs a simple aggregation:

SELECT id, SUM(gmv) FROM some_table GROUP BY id

This will produce a DataFusion execution plan that looks something like

CoalesceBatchesExec: target_batch_size=4096 RepartitionExec: partitioning=Hash([Column { name: “id”, index: 0 }], 4) AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[SUM(some_table.gmv)] TableScan: some_table

The kapot DistributedPlanner will turn this into a distributed plan by creating a shuffle boundary (called a “Stage”) whenever the underlying plan needs to perform a repartition. In this case we end up with a distributed plan with two stages:

ExecutionGraph[job_id=job, session_id=session, available_tasks=1, complete=false] =========UnResolvedStage[id=2, children=1]========= Inputs{1: StageOutput { partition_locations: {}, complete: false }} ShuffleWriterExec: None AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[SUM(?table?.gmv)] CoalesceBatchesExec: target_batch_size=4096 UnresolvedShuffleExec =========ResolvedStage[id=1, partitions=1]========= ShuffleWriterExec: Some(Hash([Column { name: “id”, index: 0 }], 4)) AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[SUM(?table?.gmv)] TableScan: some_table

The DAG structure of this ExecutionGraph is encoded in the stages. Each stage’s input field will indicate which stages it depends on, and each stage’s output_links will indicate which stage it needs to publish its output to.

If a stage has output_links is empty then it is the final stage in this query, and it should publish its outputs to the ExecutionGraphs output_locations representing the final query results.

Implementations

impl ExecutionGraph

pub fn new( scheduler_id: &str, job_id: &str, job_name: &str, session_id: &str, plan: Arc<dyn ExecutionPlan>, queued_at: u64, ) -> Result<Self>

pub fn job_id(&self) -> &str

pub fn job_name(&self) -> &str

pub fn session_id(&self) -> &str

pub fn status(&self) -> &JobStatus

pub fn start_time(&self) -> u64

pub fn end_time(&self) -> u64

pub fn stage_count(&self) -> usize

pub fn next_task_id(&mut self) -> usize

pub fn is_successful(&self) -> bool

An ExecutionGraph is successful if all its stages are successful

pub fn is_complete(&self) -> bool

pub fn revive(&mut self) -> bool

Revive the execution graph by converting the resolved stages to running stages If any stages are converted, return true; else false.

pub fn update_task_status( &mut self, executor: &ExecutorMetadata, task_statuses: Vec<TaskStatus>, max_task_failures: usize, max_stage_failures: usize, ) -> Result<Vec<QueryStageSchedulerEvent>>

Update task statuses and task metrics in the graph. This will also push shuffle partitions to their respective shuffle read stages.

pub fn running_stages(&self) -> Vec<usize>

Return all the currently running stage ids

pub fn running_tasks(&self) -> Vec<RunningTaskInfo>

Return all currently running tasks along with the executor ID on which they are assigned

pub fn available_tasks(&self) -> usize

Total number of tasks in this plan that are ready for scheduling

pub fn pop_next_task( &mut self, executor_id: &str, ) -> Result<Option<TaskDescription>>

Get next task that can be assigned to the given executor. This method should only be called when the resulting task is immediately being launched as the status will be set to Running and it will not be available to the scheduler. If the task is not launched the status must be reset to allow the task to be scheduled elsewhere.

pub fn update_status(&mut self, status: JobStatus)

pub fn output_locations(&self) -> Vec<PartitionLocation>

pub fn reset_stages_on_lost_executor( &mut self, executor_id: &str, ) -> Result<(HashSet<usize>, Vec<RunningTaskInfo>)>

Reset running and successful stages on a given executor This will first check the unresolved/resolved/running stages and reset the running tasks and successful tasks. Then it will check the successful stage and whether there are running parent stages need to read shuffle from it. If yes, reset the successful tasks and roll back the resolved shuffle recursively.

Returns the reset stage ids and running tasks should be killed

pub fn resolve_stage(&mut self, stage_id: usize) -> Result<bool>

Convert unresolved stage to be resolved

pub fn succeed_stage(&mut self, stage_id: usize) -> bool

Convert running stage to be successful

pub fn fail_stage(&mut self, stage_id: usize, err_msg: String) -> bool

Convert running stage to be failed

pub fn rollback_running_stage( &mut self, stage_id: usize, failure_reasons: HashSet<String>, ) -> Result<Vec<RunningTaskInfo>>

Convert running stage to be unresolved, Returns a Vec of RunningTaskInfo for running tasks in this stage.

pub fn rollback_resolved_stage(&mut self, stage_id: usize) -> Result<bool>

Convert resolved stage to be unresolved

pub fn rerun_successful_stage(&mut self, stage_id: usize) -> bool

Convert successful stage to be running

pub fn fail_job(&mut self, error: String)

fail job with error message

pub fn succeed_job(&mut self) -> Result<()>

Mark the job success

Trait Implementations

impl Clone for ExecutionGraph

fn clone(&self) -> ExecutionGraph

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ExecutionGraph

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl From<&ExecutionGraph> for JobOverview

fn from(value: &ExecutionGraph) -> Self

Converts to this type from the input type.