[PoC/Proposal] AQE-lite: change plan properties at runtime based on stats from pipeline breakers

[avantgardnerio]

Status: PoC/proposal, not for merge. Drafted to frame the design and get feedback from the community before any larger investment.

What the framework produces (concrete)

A query like SELECT bg.group_key, bg.sum_payload, s.payload FROM (SELECT group_key, SUM(payload) FROM big GROUP BY group_key) bg JOIN small s ON bg.group_key = s.id produces this instrumented physical plan:

01)RuntimeOptimizerExec
02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, group_key@0)], projection=[group_key@2, sum_payload@3, payload@1]
03)----CoalescePartitionsExec
04)------DataSourceExec: partitions=4, partition_sizes=[1, 0, 0, 0]
05)----ProjectionExec: expr=[group_key@0 as group_key, sum(big.payload)@1 as sum_payload]
06)------PipelineBreakerBuffer
07)--------AggregateExec: mode=FinalPartitioned, gby=[group_key@0 as group_key], aggr=[sum(big.payload)]
08)----------RepartitionExec: partitioning=Hash([group_key@0], 4), input_partitions=4
09)------------PipelineBreakerBuffer
10)--------------AggregateExec: mode=Partial, gby=[group_key@0 as group_key], aggr=[sum(big.payload)]
11)----------------DataSourceExec: partitions=4, partition_sizes=[4, 3, 3, 3]

• RuntimeOptimizerExec at line 01 — coordinator inserted at the plan root.
• PipelineBreakerBuffer at lines 06 and 09 — synchronization wrappers above each pipeline breaker (AggregateExec).
• Static planner has picked the wrong build side: CollectLeft with small (100 rows, line 03-04) as build, aggregated_big (5 actual rows, line 05-11) as probe. The runtime stat from line 07's AggregateExec(FinalPartitioned) proves this inversion. The example RuntimeRule detects it and logs swap intent:

SwapBuildSideIfInverted: would flip HashJoinExec — current build (left) = 100 rows, probe (right) = 5 rows.

Motivation

DataFusion lacks a story for adaptive query execution — runtime-informed re-optimization based on stats only knowable mid-execution (true post-aggregate cardinality, sort extrema, actual join input sizes, partition skew, etc.). Spark added AQE in 3.0 by materializing shuffle outputs to disk between stages and re-entering the optimizer. That model is a poor fit for DataFusion because DataFusion is in-process and streaming; materializing intermediate results between stages would impose Spark-style overhead on a system that doesn't have Spark-style network shuffle costs to amortize.

This PR proposes a streaming-native AQE model: insert lightweight synchronization wrappers above pipeline-breaking operators; use a shared coordinator wrapper at the plan root that runs rules at the natural barrier points; rules observe runtime stats from already-materialized in-memory state and mutate adaptive operators in place.

No materialization to disk. No re-planning pass. No extra buffering beyond what pipeline breakers already do internally. (edit: we probably could re-plan non-executed-yet plan nodes in RuntimeOptimizerExec if we want to go that route - only the interior of the plan tree would change, opaque to observers)

Components

PipelineBreakerBuffer

Sync wrapper inserted above each pipeline-breaking operator (currently AggregateExec, SortExec). Holds one batch per input partition, signals is_ready when all partitions have produced their first poll result (or terminated empty). State machine:

NeedFirstBatch  →  WaitForStreaming  →  EmitHeld  →  Streaming
                       (held)            (release)    (forward)

Three flags:

• is_ready (mechanical): every input partition has produced its first poll result.
• streaming_enabled (rule-controllable proposal): reset each poll cycle by RTO — set to is_ready as the permissive default. Rules can flip to false to veto.
• streaming_started (actual emission control): only RTO flips this via start_streaming(), which also wakes per-partition wakers.

RuntimeOptimizerExec

Inserted at the plan root by a new InsertRuntimeOptimizer physical optimizer rule. On every poll_next, runs a three-phase coordinator cycle:

1. Propose: walk the subtree; set streaming_enabled = is_ready on each buffer (permissive default).
2. Evaluate: run each registered RuntimeRule. Rules may flip streaming_enabled = false (veto), mutate adaptive operators in place, or both.
3. Commit: walk the subtree; for each buffer still enabled, call start_streaming().

A shared Arc<AtomicWaker> is threaded into both RTO and every buffer. Buffers wake it when their is_ready flips; RTO registers cx.waker() on it before each walk. The AtomicWaker is essential because cx.waker() is task-local — a buffer's poll inside a spawned subtask (e.g. one of RepartitionExec's internals) cannot wake RTO via cx.waker(), but can via the shared AtomicWaker.

RuntimeRule trait

pub trait RuntimeRule: Send + Sync + std::fmt::Debug {
    fn name(&self) -> &str;
    fn evaluate(&self, plan: &Arc<dyn ExecutionPlan>);
}

Cheap, idempotent visitors. Rules track their own "already fired" state. They observe runtime stats via new trait methods like runtime_row_count (this PR) and (future) runtime_partition_extrema (analogous to #23090), and mutate adaptive operators via typed methods (HashJoinExec::flip_sides, RepartitionExec::set_split_points, etc. — all future work).

ExecutionPlan::runtime_row_count(partition)

New trait method (default None) returning the number of rows an operator will emit on a given partition. For pipeline-breaking operators, the value is knowable the moment their build phase completes — before any output batch is pulled. Implemented on AggregateExec here, reading from the existing OutputRows metric. ProjectionExec and PipelineBreakerBuffer are passthroughs (the buffer's passthrough is gated: only delegates once is_ready).

Example rule: SwapBuildSideIfInverted

First concrete RuntimeRule. Walks the plan looking for HashJoinExec. For each, compares its left (build) and right (probe) sides' row counts: static statistics() when Exact (e.g. an in-memory dimension table), runtime via runtime_row_count summed across partitions otherwise (e.g. the output of a GROUP BY whose distinct cardinality the planner couldn't predict).

If the build side is larger than the probe side, the rule would call HashJoinExec::flip_sides(). But that method doesn't exist yet, so this PR logs intent only. The detection mechanism is fully wired end-to-end; the actually-saving-work part needs HashJoin to defer its build phase until both sides' pipeline-breakers report is_ready — a follow-up operator change.

SLT walkthrough

runtime_optimizer.slt sets up the canonical misestimation case: big (100K rows) joined to small (100 rows) through a GROUP BY on a low-cardinality column the planner can't propagate distinct stats through. JoinSelection trusts the Inexact ~100K estimate for aggregated_big and picks small as the build side. At runtime aggregated_big is only 5 distinct groups; the rule detects the inversion and would flip.

The EXPLAIN block in the SLT shows the static plan with the AQE primitives wrapped — RuntimeOptimizerExec at root, PipelineBreakerBuffer above each AggregateExec. The build-side swap is not reflected in EXPLAIN because the proposed runtime mutation cannot retroactively appear in static plan output.

What lands here vs. what's deferred

Lands:

• PipelineBreakerBuffer operator + is_pipeline_breaker heuristic + insertion rule.
• RuntimeOptimizerExec operator + insertion rule + shared AtomicWaker wake propagation.
• RuntimeRule trait + three-phase coordinator (propose / evaluate / commit).
• ExecutionPlan::runtime_row_count trait method.
• AggregateExec impl of runtime_row_count; ProjectionExec and PipelineBreakerBuffer passthroughs.
• SwapBuildSideIfInverted rule (log-only).
• SLT demonstrating the detection end-to-end.

Deferred to follow-ups (each its own PR-sized piece of work):

1. HashJoinExec::flip_sides() with interior mutability + a build-phase deferral mechanism that waits for descendant PipelineBreakerBuffers to be ready before starting build. This is the change that turns SwapBuildSideIfInverted from a logger into a real optimization.
2. RepartitionExec::set_split_points() — mutable Partitioning::Range(Option<Vec<SplitPoint>>). Pairs with runtime_partition_extrema on SortExec to enable parallel cumulative-window functions and parallel range repartitioning in general. Mutating RepartitionExec is timing-safe (it has no synchronous build phase, just routing), so this follow-up doesn't need operator deferral.
3. Additional rules: partition coalescing, skew splitting, adaptive aggregation strategy, broadcast-join switching.
4. runtime_partition_extrema trait method on SortExec, parallel-window operators (CarryExec, HaloDropExec) — there's already a parallel PoC at Parallel bounded RANGE-frame window functions without PARTITION BY (draft) #23026 / feat: PartitionExtrema + SortExec observer + BoundedWindowAggExec passthrough #23090 / feat(physical-expr): add Partitioning::DynamicRange variant #23094 that overlaps significantly; this work would either build on or replace those primitives.
5. Restore is_pipeline_breaker() trait method. Pre-Replace execution_mode with emission_type and boundedness #13823 there was a direct API for this; after the split into EmissionType + Boundedness it has to be reconstructed. Neither EmissionType::Final (inherited from descendants, so wraps too aggressively) nor the transition-point filter (Final && all-children-non-Final, misses cascading breakers like FinalPartitioned above Partial) is right. The principled fix is a dedicated method; we use a hardcoded match for now.

Discussion points I'd appreciate feedback on

• Architectural model. Is the propose/evaluate/commit cycle the right shape, or should rules return decisions to RTO instead of mutating buffer state directly?
• Pipeline-breaker detection. See deferred point 5 — happy to upstream a dedicated is_pipeline_breaker() trait method as a small precursor PR if there's appetite.
• Trait surface. runtime_row_count(partition: usize) returns total per partition. Is per-partition vs. cross-partition total the right primitive?
• Wake propagation via AtomicWaker. Is the shared waker pattern acceptable, or is there a more idiomatic approach DataFusion already uses for cross-task notification I'm missing?
• Naming. RuntimeOptimizerExec vs AdaptiveExec vs something else? Open to bikeshedding.

What this PoC is not

• Not a complete AQE implementation. The log-only rule means no actual work is saved yet.
• Not a replacement for the existing parallel-window-poc work — that work targets the same general space (runtime-informed plan adaptation) but via different primitives (the runtime_partition_extrema API + new operators). The two designs are compatible and could converge.
• Not optimized. The on-every-poll subtree walks are O(plan depth); fine for the prototype, would want amortization in production.

Looking for: "is this the right architectural direction for AQE in DataFusion?" before investing in any of the follow-up work.

[thinkharderdev]

They observe runtime stats via new trait methods like runtime_row_count (this PR) and (future) runtime_partition_extrema (analogous to #23090)

Is there a reason we can't do this using the existing ExecutionPlan::partition_statistics method?

[avantgardnerio]

using the existing ExecutionPlan::partition_statistics

Yes, that's the case above. The static statistics lie, because they can't know the cardinality reduction. The aggregate does before it emits a batch, thus the flip from the static plan in the dynamic one.

Edit: if you mean use the same method to return runtime statistics - possibly, yes. It has pass-through implications though for children.

[comphead]

Thanks @avantgardnerio this is challenging task, some insights how Spark AQE does the same

Spark plans a shuffle join initially, lets the input shuffle map stages actually run, reads the byte counts that came out of those stages, and then re-runs the optimizer + planner on the still-unexecuted portion of the plan with those real numbers in hand. Because the size is now real instead of estimated, the planner's normal JoinSelection strategy can flip a SortMergeJoinExec into a BroadcastHashJoinExec — the rule that picks broadcast is the same rule the non-AQE planner uses; what's different is the stats it sees.

Spark is able to identify if relationship is large or small based on the shuffle size, data which available in runtime before the join.

For proposed approach I'm probably missing the part to identify size of relationship? if the first batch is a small batch and subsequent are large ones, what decision would be made?

@thinkharderdev Reg to stats, are they calculated on runtime? if one of the relationship is filtered/generated table, how rowcounts calculated for partition stats?

[avantgardnerio]

if the first batch is a small batch and subsequent are large ones

The idea is that the PipelineBreakerBuffer only comes after pipeline breakers - which consume all input batches before emitting 1 output batch; natural stage barriers within our existing plans. By the time PipelineBreakerBuffer::is_ready() == true, we know the exact stats from the leaf, not a 1 batch approximation.

[avantgardnerio]

Spark plans a shuffle join initially, lets the input shuffle map stages actually run, reads the byte counts that came out of those stages, and then re-runs the optimizer...

Spark's shuffle is a pipeline breaker; this PR is stating that any pipeline breaker gives us the same epistemic guarantee, just at finer granularity than full-stage shuffles

[comphead]

Spark's shuffle is a pipeline breaker; this PR is stating that any pipeline breaker gives us the same epistemic guarantee, just at finer granularity than full-stage shuffles

Shuffle is not always a pipeline breaker, but the concept of pipeline breaker would help in DF if we want to deal with runtime decision, to consume all inputs before producing output. Overall the design makes sense to me IMO. But how would you use it for joins anyway? 🤔

For SMJ the sorting is a pipeline breaker.
What about HJ, NLJ? What would be pipeline breakers for them?

[thinkharderdev]

Reg to stats, are they calculated on runtime? if one of the relationship is filtered/generated table, how rowcounts calculated for partition stats?

What I am imagining is that the operators update stats as they process the data. So pre-execution the operator returns partition stats based on static statistics (like it does now) but as data flows through they are updated. For a pipeline breaker if you call partition_statistics once all input is consumed you get back Precision::Exact stats because the precise values are known.

I'm concerned that the approach outlined in this PoC is moving towards an entirely separate and parallel optimization framework which seems confusing. There shouldn't really be a difference between optimizing a plan pre-execution vs in-flight as ultimately you are just (in principle) doing the same optimization passes but with better statistics.

[Dandandan]

I agree with the points already made above:

• It should be clear when a non-shuffle/non-push-based execution will benefit from this: for joins, we usually don't want to fully know both sides of the join as they need to be spilled / shuffled fully when the underlying data source / query to get exact statistics. AQE in streaming mode will be different than when data is streamed / pushed.

• We probably want to reuse the existing optimizations as much as possible. E.g. JoinSelection etc.

Also I am a bit confused to see RuntimeOptimizerExec, I think that feels a bit hacky to reuse the execution plan for some AQE plumbing 🤔 (but perhaps I need to be convinced)

[avantgardnerio]

Thanks for the feedback guys! I really appreciate all of it.

I think I can answer most of it in a new push, coming shortly. Once that's up I'll address each point.

[avantgardnerio]

Thanks again for the feedback! I think it has helped me shape this into a much more AQE-looking PR:

0. The PR is now functional. It actually does flip the join sides. The SLT passes just like it did without AQE. Based on runtime_row_count()
1. PipelineBreakerBuffer is now StageBoundryBuffer, and it actually buffers. This allows us to insert it even on the left (streaming side) of the example SLT HashJoin.
2. Each StageBoundry gets assigned a number, within a stage all leaves execute in parallel. Stages execute sequentially.
3. RuntimeOptimizer calls prime() on each StageBoundryBuffer to kick it off in it's own tokio task, much like how RepartitionExec polls it's children today - this skips the normal DataFusion poll-based flow to execute subtrees incrementally and bottom up
4. RuntimeRule now has the exact same signature as PhysicalOptimizerRule, hopefully showing how they can be the same (just requires refactoring into a common crate). Rules would have to be made "runtime aware" one at a time, but this gives us a migration path to AQE
5. StageBoundryBuffer is it's own pipeline breaker - it stores everything in an unbounded channel in this PR. I think in the future we could prime() it until it hits a memory threshold, then give up and release it to start streaming - if this happens we'd no longer be able to get run time stats and modify the plan, but at least it would be "no worse" than today.

  RTO: stage 0 ready (2 boundaries); firing 1 rule(s) before release
  SwapBuildSideIfInverted: flipping HashJoinExec — current build (left) = 100 rows, probe (right) = 5 rows.
  RTO: stage 0 released; downstream consumers can now drain ...

01)RuntimeOptimizerExec
02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, group_key@0)], projection=[group_key@2, sum_payload@3, payload@1]
03)----StageBoundaryBuffer: stage=0
04)------CoalescePartitionsExec
05)--------DataSourceExec: partitions=4, partition_sizes=[1, 0, 0, 0]
06)----StageBoundaryBuffer: stage=0
07)------ProjectionExec: expr=[group_key@0 as group_key, sum(big.payload)@1 as sum_payload]
08)--------AggregateExec: mode=FinalPartitioned, gby=[group_key@0 as group_key], aggr=[sum(big.payload)]
09)----------RepartitionExec: partitioning=Hash([group_key@0], 4), input_partitions=4
10)------------AggregateExec: mode=Partial, gby=[group_key@0 as group_key], aggr=[sum(big.payload)]
11)--------------DataSourceExec: partitions=4, partition_sizes=[4, 3, 3, 3]

Hopefully this makes clear that the PR is a migration path, a generalizable solution. So what's missing? If we keep going this route would we be able to 1. land it today, 2. eventually get to "full" AQE?

[avantgardnerio]

And as promised, answers to individual questions:

Why not extend the existing partition_statistics method instead of adding runtime_row_count?

I think we should do that. What I have was convenient for fall through to children.

Feels like a parallel optimization framework. There shouldn't be a difference between pre-execution and in-flight optimization.

That's exactly the target, and it's why RuntimeRule has the same signature as PhysicalOptimizerRule now. The two traits exist today only because physical-plan can't depend on physical-optimizer (the dependency runs the other way) . End-state: one trait, with rules optionally reading runtime stats from boundaries in the plan tree they receive. Static rules ignore runtime state; runtime-aware rules use it. Per-rule migration: any existing PhysicalOptimizerRule (including JoinSelection) can be made runtime-aware by adding state inspection. Existing rules don't change until they want to.

Pipeline breakers for HJ and NLJ?

For HJ this PR inserts one: InsertHashJoinBoundaries wraps each HashJoin input in a StageBoundaryBuffer. The buffer materializes its input and IS the pipeline breaker — we don't need a natural one inside the join.

• Same pattern works for NLJ (no natural breaker, just insert one above each input — ~30-line rule).
• For SMJ either approach works: reuse the existing SortExec (it's already a breaker — insert the boundary above it for stats) or insert directly above the SMJ inputs. Either way the infrastructure is the same.
• The contribution is precisely this generalization: any operator can become adaptive without depending on a natural breaker existing inside it.

Joins usually don't want to fully know both sides; when does non-shuffle execution benefit?

• The win case is exactly what the SLT shows: one side is derived (aggregation, filter, scan-with-projection) whose post-operator size is much smaller than the planner's Inexact upper bound. Materializing that side is cheap relative to the bad build-side choice it lets us avoid.
• For sides where materialization is genuinely unsafe (large streaming inputs), the stream-through fallback (point 5 of the summary) gates: if memory pressure hits the threshold, the boundary releases without finishing materialization, the runtime rule sees an unflagged stage and bails, the query runs exactly as it would today.
• Net: adaptive wins when stats are knowable cheaply; otherwise we're no worse than the pre-AQE baseline.

Reuse existing optimizations like JoinSelection.

• Strongly agree. SwapBuildSideIfInverted already calls HashJoinExec::swap_inputs — the same primitive JoinSelection uses. Same logic, fresher numbers.
• With the trait unification above, JoinSelection itself could run at runtime against the post-breaker plan, no re-implementation needed.

RuntimeOptimizerExec feels hacky — reusing an ExecutionPlan for AQE plumbing.

• Fair concern. The reason it IS an ExecutionPlan: that's the only point in the existing pipeline where we can intercept poll-time decisions without changes to every operator or to the session layer. The RTO sits at the root, opt-in via an optimizer rule.
• The wrapper-at-root pattern matches other "coordinator" nodes in DF (e.g. how CoalescePartitionsExec orchestrates downstream). The operator interface is the natural composition surface.
• Open to alternatives — a SessionContext-level coordinator would also work but requires plumbing into every execute path. The wrapper was the surgical "get it done today, show working code" compromise of this PR.

Summary: RuntimeRule collapses into PhysicalOptimizerRule, runtime_row_count collapses into partition_statistics, RuntimeOptimizerExec becomes the only new operator, and StageBoundaryBuffer is the only new infrastructure — everything else is existing rules reading more accurate stats.