[EPIC] Adaptive Query Execution (AQE)

Executive Summary

Create a distributed query planner which can change stages, add/or remove stages, or cancel stage (partitioning) to reflect collected runtime statistics.
An extendable distributed query planner, simple to extend, just like adding a new physical rule.

Overview

Currently, jobs in Ballista are divided into stages at job submission time according to predefined rules.
Ballista can optimize stages based on runtime information using set of physical optimizers, changing the join order for example. This has been implemented here with only few physical optimizer rules used.

Also, few more rules are scattered around here and here

Current approach may fail to generate valid plan in case there is dependency between physical optimizers. Furthermore, some optimizations can’t be propagated across stage boundaries potentially impacting overall performance.

Physical optimizers used work well with a major downside, rules being scattered across code, making adding or removing rules non-trivial task. Scattered rules are hard to control, thus centralized rule "facility" is need. It would make it easier to change or alter behavior of physical planners or to experiment with new solutions.

Another major problem arises when stages are split; it becomes impossible to change their layout, which can be necessary in certain situations. For instance, if an exchange produces an empty result, we’d need to run a subsequent stage even though the system knows it will also produce an empty result. This can lead to significant overhead on the scheduler, especially with many tasks. Having the ability to remove a stage would be helpful in this scenario.

Proposal

Implement adaptive query planner which would address some of the issues with the current approach. In the core of the design are pluggable physical optimizer rules which would re-optimize plan after each stage.

This proposal is strongly influenced by Adaptive and Robust Query Execution for Lakehouses at Scale paper.

Major difference from it is that ballista works on physical plan rather than logical plan.

Impact on DataFusion

Currently, in DataFusion, physical optimizer rules are run once per plan execution, with current proposal, whole set of rules will be run after each stage, thus if physical optimizer rules are not idempotent, they may keep adding physical plane exec nodes even they have been already added and unnecessary.

Physical optimizer rules have single pass, which may not work well with current approach (see "adding new stage" after join reordering further below).

Impact on Ballista

Stage planning in Ballista revolves around ExecutionGraph, which prescribes a translation from a physical plan to a set of stages executing that plan. ExecutionGraph was implemented as a concrete trait, hence it’s not a trivial task to replace it. ExecutionGraph does accept a dyn DistributedPlanner, but unfortunately DistributedPlanner returns a static list of stages, which would not work with adaptive query planning.

ExecutionGraph is initialized in TaskManager when new job is submitted.

There are two approaches which we can take with ExecutionGraph:

Keep ExecutionGraph as it is, adding code required for AQE. This approach might be possible, but it would be hard to maintain, and probably messy to implement.
Rename ExecutionGraph to StaticExecutionGraph introduce dyn ExecutionGraph which would be implemented by StaticExecutionGraph and new AdaptiveExecutionGraph. Change implementation in TaskManger based on predefined configuration parameter. This approach looks like easier route to take. It would enable us to have two implementation running in parallel until AQE implementation matures.

Adaptive planner will need some time to get it right. It will work with current distributed planner for foreseeable future, will be disabled by default with configuration option to turn it on.

Plan Arithmetics

Adding New Stage to Running Plan

Some optimization may produce a need for an exchange to be inserted into a running plan bug 1 and bug 2

┌───────────────────────────┐
│   AdaptiveDatafusionExec  │
│    --------------------   │
│      is_final: false      │
│         plan_id: 1        │
│                           │
│          stage_id         │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│       CrossJoinExec       ├──────────────┐
└─────────────┬─────────────┘              │
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│   CoalescePartitionsExec  ││      CooperativeExec      │
│                           ││    --------------------   │
│                           ││      CooperativeExec      │
└─────────────┬─────────────┘└─────────────┬─────────────┘
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│        ExchangeExec       ││         File Scan         │
│    --------------------   ││    --------------------   │
│       partitioning:       ││     num_partitions: 2     │
│    Hash([big_col@0], 2)   ││                           │
│                           ││        statistics:        │
│         plan_id: 0        ││ [Rows=Inexact(1024), Bytes│
│       stage_id: None      ││ =Inexact(8192), [(Col[0]: │
│   stage_resolved: false   ││            )]]            │
└─────────────┬─────────────┘└───────────────────────────┘
┌─────────────┴─────────────┐
│      CooperativeExec      │
│    --------------------   │
│      CooperativeExec      │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         File Scan         │
└───────────────────────────┘

plan after join reordering and new stage addition (plan_id: 3), looks similar to:

┌───────────────────────────┐
│   AdaptiveDatafusionExec  │
│    --------------------   │
│       is_final: true      │
│         plan_id: 1        │
│        stage_id: 2        │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│       ProjectionExec      │
│    --------------------   │
│      big_col: big_col     │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│       CrossJoinExec       ├──────────────┐
└─────────────┬─────────────┘              │
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│   CoalescePartitionsExec  ││        ExchangeExec       │
│                           ││    --------------------   │
│                           ││       partitioning:       │
│                           ││    Hash([big_col@0], 2)   │
│                           ││                           │
│                           ││         plan_id: 0        │
│                           ││     stage_id: 0           │
│                           ││    stage_resolved: true   │
└─────────────┬─────────────┘└─────────────┬─────────────┘
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│        ExchangeExec       ││      CooperativeExec      │
│    --------------------   ││    --------------------   │
│       partitioning:       ││      CooperativeExec      │
│           None            ││                           │
│                           ││                           │
│         plan_id: 3        ││                           │
│     stage_id: 1           ││                           │
│    stage_resolved: true   ││                           │
└─────────────┬─────────────┘└─────────────┬─────────────┘
┌─────────────┴─────────────┐┌─────────────┴─────────────┐
│      CooperativeExec      ││         File Scan         │
│    --------------------   ││                           │
│      CooperativeExec      ││                           │
└─────────────┬─────────────┘└───────────────────────────┘
┌─────────────┴─────────────┐
│         File Scan         │
│    --------------------   │
│     num_partitions: 2     │
│                           │
│        statistics:        │
│ [Rows=Inexact(1024), Bytes│
│ =Inexact(8192), [(Col[0]: │
│            )]]            │
└───────────────────────────┘

Removing Stages from a Running Plan

Apart from adding stages, stages should be removed if they will produce empty result.
This will reduce pressure on scheduler as it will skip task scheduling.

Take for example plan

[stage 2] --------------------------------------------------------------------------------------------------
AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
  ProjectionExec: expr=[c0@0 as c0, count(Int64(1))@1 as count(*)]
    AggregateExec: mode=FinalPartitioned, gby=[c0@0 as c0], aggr=[count(Int64(1))]
      CoalesceBatchesExec: target_batch_size=8192
       [stage 1] -------------------------------------------------------------------------------------------      
        ExchangeExec: partitioning=Hash([c0@0], 2), plan_id=1, stage_id=pending, stage_resolved=false
          AggregateExec: mode=Partial, gby=[c0@0 as c0], aggr=[count(Int64(1))]
            ProjectionExec: expr=[min(t.a)@1 as c0]
              AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a)]
                CoalesceBatchesExec: target_batch_size=8192
                 [stage 0] ---------------------------------------------------------------------------------                
                  ExchangeExec: partitioning=Hash([c@0], 2), plan_id=0, stage_id=pending, stage_resolved=false
                    AggregateExec: mode=Partial, gby=[c@1 as c], aggr=[min(t.a)]
                      CoalesceBatchesExec: target_batch_size=8192
                        FilterExec: a@0 = 42
                          DataSourceExec: partitions=1, partition_sizes=[1]

Running first stage, producing no results

ShuffleWriterExec: partitioning:Hash([c@0], 2)
  AggregateExec: mode=Partial, gby=[c@1 as c], aggr=[min(t.a)]
    CoalesceBatchesExec: target_batch_size=8192
      FilterExec: a@0 = 42
        DataSourceExec: partitions=1, partition_sizes=[1]

statistic information will be used, by adaptive planner to skip all other stages and just return empty result.

┌───────────────────────────┐
│     ShuffleWriterExec     │
│    --------------------   │
│     partitioning: None    │
└─────────────┬─────────────┘
┌─────────────┴─────────────┐
│         EmptyExec         │
└───────────────────────────┘

Coalesce Shuffle Partitions (Elastic Shuffle Parallelism)

Should shuffle partitions produce very little data, planner should coalesce partitions, reducing number of task needed to produce results. The rule modifies the partitioning specification in the ShuffleRead operator

Canceling Redundant (Running) Stages

Adding or removing stages may arise need to cancel already running stages. This can happen, for example, if some stage produces empty result making other stages unnecessary, or if there is better, semantic equal but considered superior plan (this case needs further investigation).

Take for example plan

AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
  CoalesceBatchesExec: target_batch_size=8192
    HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0)], projection=[a@1, b@2, c@3]
      CoalesceBatchesExec: target_batch_size=8192
        ExchangeExec: partitioning=Hash([a@0], 2), plan_id=0, stage_id=0, stage_resolved=false
          CoalesceBatchesExec: target_batch_size=8192
            FilterExec: b@1 = 42, projection=[a@0]
              DataSourceExec: partitions=1, partition_sizes=[1]
      CoalesceBatchesExec: target_batch_size=8192
        ExchangeExec: partitioning=Hash([a@0], 2), plan_id=1, stage_id=1, stage_resolved=false
          DataSourceExec: partitions=1, partition_sizes=[1]

we have two runnable stages, a build side of join ExchangeExec: partitioning=Hash([a@0], 2), plan_id=0, stage_id=0, stage_resolved=false and a probe side ExchangeExec: partitioning=Hash([a@0], 2), plan_id=1, stage_id=1, stage_resolved=false.

Let's say probe side needs some time to finish.

If the build side, stage (0), finishes first with shuffle producing no data hash join will produce no data as well, thus whole hash join could be replaced with EmptyExec.
Something like:

AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending
  EmptyExec

Plan could not finish until both stages finish, increasing overall result latency. Cancelling redundant stages would make sense, as it will reduce cluster utilization (and result latency).

NOTE: current "plan splitter" implementation draft https://github.com/milenkovicm/datafusion-ballista/blob/707ad0bcdf38e6aff02841b4ea3afe3ad7833264/ballista/scheduler/src/state/aqe/mod.rs#L60

current tasks

feat: Extract execution_graph to a trait #1361
feat: Adaptive query execution (AQE) planner fundamentals #1372
join reordering feat: Adaptive query execution (AQE) planner fundamentals #1372
elimination of stages which do not produce data (initial implementation) feat: Adaptive query execution (AQE) planner fundamentals #1372
Ensure idempotency in the DataFusion physical optimizer to support Ballista AQE #1378
verify if we can reuse instead of creating our own rule for eliminating joins
support executor&stage failure
feat(aqe): CoalescePartitionsRule — shuffle-partition coalescing on resolved stats #1684
feat(aqe): delay join decision & introduce broadcast join in AQE #1752
switch switch from streaming aggregation to hash aggregation (extended rules)
planner traverses whole job graph every time, even for stages which are completed (can we trim the graph and keep traversal faster?)
support sort shuffle

@andygrove

I plan to put more time on this functionality in coming weeks (hopefully), would @andygrove , @alamb or @Dandandan have some inputs.

I will try and review this tomorrow

@mattcuento

Implement adaptive query planner which would address some of the issues with the current approach. In the core of the design are pluggable physical optimizer rules which would re-optimize plan after each stage.

This is my understand of current best practice with distributed engines -- that is be able to adaptively re-organize after the plan has started. Given that Ballista already breaks the execution into stages where the intermediate results are persisted, I think this is

Rename ExecutionGraph to StaticExecutionGraph introduce dyn ExecutionGraph which would be implemented by StaticExecutionGraph and new AdaptiveExecutionGraph. Change implementation in TaskManger based on predefined configuration parameter. This approach looks like easier route to take. It would enable us to have two implementation running in parallel until AQE implementation matures.
Adaptive planner will need some time to get it right. It will work with current distributed planner for foreseeable future, will be disabled by default with configuration option to turn it on.

This also makes sense to me as it will allow a phased rollout of the new adaptive execution rather than having to get it all right at once. Eventually you might be able to use the dynamic planner by default, and detect patterns in the plan that it can't handle and fall back to the static planner

Maybe @mattcuento or @killzoner have some additional thoughts

@alamb

Thanks @alamb!
I was hoping @killzoner, @mattcuento, @danielhumanmod might be interested contributing some part of it

@milenkovicm

Thanks @alamb! I was hoping @killzoner, @mattcuento, @danielhumanmod may be interested contributing some part of it

Thank you @milenkovicm for the detailed write-up!

I took a look at your PRs (#1361 and #1372). It looks like the core ExecutionGraph refactoring and the state transition logic (revive/rollback/update) are coming along nicely. Appreciate your work on building the foundation.

Do you have a specific task or area you would recommend I start with? I'm happy to pick up whatever fits best with your current progress and priorities.

@danielhumanmod

thanks @danielhumanmod, unless you wish to continue work on #1372 i dont have much at the moment, as makes no sense opening other tasks until we get core working.

one thing which could be fixed, some of the physical optimiser rules are not idempotent (every time called they add nodes), please have a look at https://github.com/milenkovicm/datafusion-ballista/blob/10fb25b9836e70cca63e86728fbf4ebee3e257c3/ballista/scheduler/src/state/aqe/optimizer_rule/datafusion_patch.rs it might be useful to fix those in datafusion

@milenkovicm

Thanks @alamb! I was hoping @killzoner, @mattcuento, @danielhumanmod might be interested contributing some part of it

Hey @milenkovicm, happy to work on this. Of the two approaches described, i'm also feeling like the 2nd one is easier both for testing and landing changes gradually, rather than a big one time change.

Also thank you for all the details on inner machinery, it's very valuable for a newcomer like me. Appreciate any guideline as I still have a very incomplete knowledge of the codebase and inner details at the time. Should we drop other stuff for the moment in favor of this one?

@killzoner

thanks @killzoner
please do not drop other things you're working on just yet. I'll try to get into shape candidate branch first, unfortunately its going slower than i anticipated

@milenkovic

Hey @milenkovic, also happy to help here!

A lot of the proposal makes sense to me, though I should say my experience here is limited to Flink's adaptive scheduling which is more strictly focused on parallelism.

+1 for the second approach, testability, rollout, and falling back to a static planner sound more simple in this way. Agreed that physical optimizer passes between stages with stage-based execution makes sense generally.

Two questions:

Im curious about optimizations like two-phased aggregations, which would span stages (local agg in one, global in the next). I’ve seen limited examples of adaptive two-phase aggregations where the local agg can be disabled in the presence of a large keyspace. Have you given any thought to how the design might allow/disallow future extensions for rules that span > 1 stage and aim to adapt a running stage? Sounds like such rules wouldn’t be possible with a pass after each stage, which is fine, but wanted to make that thought explicit.
For observability sake, do we plan to log out or keep in memory each round of optimizations and the original execution graph?

Im curious about optimizations like two-phased aggregations, which would span stages (local agg in one, global in the next). I’ve seen limited examples of adaptive two-phase aggregations where the local agg can be disabled in the presence of a large keyspace. Have you given any thought to how the design might allow/disallow future extensions for rules that span > 1 stage and aim to adapt a running stage? Sounds like such rules wouldn’t be possible with a pass after each stage, which is fine, but wanted to make that thought explicit.

I believe datafusion.execution.skip_partial_aggregation_probe_ratio_threshold controls this, in ballista case this would mean it will just write everything to shuffle file. Not sure if we could make some kind of optimizer rule which could use cardinality from previous stage to help with this.

Anyway, good part of work to make it easier to experiment with this things

For observability sake, do we plan to log out or keep in memory each round of optimizations and the original execution graph?

At the moment prototype holds whole plan and traverse whole plane every after every stage. It may be good for observability but it may affect planning time, thus it may not be the best way forward. At the moment we can keep it as it is until we get it running and then consider how to handle planning and observability. I do agree that is a bit of a challenge to understand what happened with plan.

Something that might be useful to include, that I think Spark-like AQE does not do, but is a super useful optimization is early stopping query execution based on global limit during a scan (e.g. select * where x limit 100).

This is not possible with just reoptimizing after a full stage or even a task has been finished (at that point other running tasks might have also already scanned way more than necessary).
This can avoid using a lot of resources / improve query latency a lot for highly selective queries.

Something that might be useful to include, that I think Spark-like AQE does not do, but is a super useful optimization is early stopping query execution based on global limit during a scan (e.g. select * where x limit 100).

This is not possible with just reoptimizing after a full stage or even a task has been finished (at that point other running tasks might have also already scanned way more than necessary). This can avoid using a lot of resources / improve query latency a lot for highly selective queries.

how can we stop stage in this case ?

Something needs to track the count of rows being read in the partitions / tasks and early

Something that might be useful to include, that I think Spark-like AQE does not do, but is a super useful optimization is early stopping query execution based on global limit during a scan (e.g. select * where x limit 100).
This is not possible with just reoptimizing after a full stage or even a task has been finished (at that point other running tasks might have also already scanned way more than necessary). This can avoid using a lot of resources / improve query latency a lot for highly selective queries.

how can we stop stage in this case ?

Something needs to track the sum of count of rows being written during scan, and cancel all the tasks/partitions (from schedulur) once the sum exceeds the limit.

@milenkovicm

Hi @milenkovicm , thanks for this great contribution/initiative. I'm wondering is there anything I can pick up? Probably I can work on #338

Maybe "support executor failure" if you wish @jja725 that one is on the critical path at the moment and I have no bandwidth at the moment.

@jja725

hey @jja725 its up to you, if you think #338 would be interesting go ahead, i can assist should you need help

Thanks I would spend sometime on it and see how far I can go

Are there any plans on introducing Spark-like DynamicJoinSelection and JoinSelection (as strategy or PhysicalOptimizerRule in terms of DF) to select between sort-merge join, hash join and broadcast hash join?
I think when #1647 is ready that would be possible.

@sandugood

Are there any plans on introducing Spark-like DynamicJoinSelection and JoinSelection (as strategy or PhysicalOptimizerRule in terms of DF) to select between sort-merge join, hash join and broadcast hash join?
I think when #1647 is ready that would be possible.

Fill free to try it @sandugood i think it make sense

I will try to tackle it

Early-stop on global LIMIT (v1) — draft implementation available at #1727
Adds scheduler-side row tracking that cancels remaining tasks once a bare LIMIT N is satisfied. Key components: JobLimitTracker (atomic one-shot trigger), LimitEarlyStopAnalyzer (plan walker for eligible GlobalLimitExec), and scheduler integration (EarlyStopCancel event + partial-success finalization). Gated behind existing AQE flag + new ballista.aqe.limit_early_stop.enabled config.

@milenkovicm

Hi @milenkovicm, I wanted to clarify the purpose of "switch from streaming agg to hash agg". Is it a correct defensive or performance improvement? If it's the correct defensive, I believe AQE already did that well. Basically, the ShuffleReaderExec advertises blank eq_properties, so the per-stage AggregateExec is auto-rebuilt as Linear (i.e. Hash). Streaming agg never runs over hash-shuffled data.

Two upstream DataFusion PRs have been opened to fix physical-optimizer rules that are not idempotent, which directly impacts AQE replanning:

These were the only two offenders found by running a discovery harness that applies each rule in PhysicalOptimizer::new().rules twice on a set of plan fixtures and compares the output structurally. Once these land in a DataFusion release, WarnOnDuplicateExecRule in Ballista's AQE chain should no longer fire on standard workloads (verified against a multi-stage aggregation fixture and TPC-H-style plans). A Ballista-side regression test is in-progress.

fix(physical-plan): make HashJoinExec dynamic filter pushdown idempot…

fix(physical-optimizer): make OutputRequirements idempotent (apache#2…

Filed discrete tracking issues for the four AQE limitations currently documented in the tuning guide / AdaptiveExecutionGraph doc comment:

Executor failure handling on the AQE path — AQE: handle executor failure on the adaptive execution path #1986
Dynamically coalesce shuffle partitions — AQE: dynamically coalesce shuffle partitions based on runtime statistics #1987
Switch hash join → sort-merge join at runtime — AQE: switch from hash join to sort-merge join based on runtime statistics #1988
Switch streaming aggregation → hash aggregation at runtime — AQE: switch from streaming aggregation to hash aggregation based on runtime statistics #1989

closing this as we have AQE enabled

Congrats everyone on this milestone!

Navigation Menu

Repository navigation

[EPIC] Adaptive Query Execution (AQE) #1359

Executive Summary

Overview

Proposal

Impact on DataFusion

Impact on Ballista

Plan Arithmetics

Adding New Stage to Running Plan

Removing Stages from a Running Plan

Coalesce Shuffle Partitions (Elastic Shuffle Parallelism)

Canceling Redundant (Running) Stages

1 remaining item

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[EPIC] Adaptive Query Execution (AQE) #1359

Description

Executive Summary

Overview

Proposal

Impact on DataFusion

Impact on Ballista

Plan Arithmetics

Adding New Stage to Running Plan

Removing Stages from a Running Plan

Coalesce Shuffle Partitions (Elastic Shuffle Parallelism)

Canceling Redundant (Running) Stages

Activity

milenkovicm commented on Jan 3, 2026

milenkovicm commented on Jan 7, 2026

alamb commented on Jan 7, 2026

alamb commented on Jan 8, 2026

milenkovicm commented on Jan 8, 2026

danielhumanmod commented on Jan 9, 2026

milenkovicm commented on Jan 9, 2026

killzoner commented on Jan 9, 2026

milenkovicm commented on Jan 9, 2026

mattcuento commented on Jan 9, 2026

milenkovicm commented on Jan 10, 2026

Dandandan commented on Jan 12, 2026

milenkovicm commented on Jan 12, 2026

Dandandan commented on Jan 13, 2026

1 remaining item

milenkovicm commented on Apr 6, 2026

milenkovicm commented on Apr 8, 2026

jja725 commented on Apr 9, 2026

sandugood commented on May 8, 2026

milenkovicm commented on May 9, 2026

sandugood commented on May 9, 2026

wirybeaver commented on May 18, 2026

wirybeaver commented on May 26, 2026

wirybeaver commented on May 27, 2026

andygrove commented on Jul 10, 2026

milenkovicm commented on Aug 29, 2026

killzoner commented on Aug 29, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions