fix: repair AQE plan rewrites, remove a redundant shuffle, and enable adaptive query planning by default - #2315
Conversation
Adaptive query planning re-optimises the physical plan after every stage completion. Four defects in that path made the TPC-DS suite hang at q23 and return wrong results for two queries. InterleaveExec requires all children to share one Hash/Range partitioning and asserts on rebuild. AQE rewrites union branches independently, and TreeNode rebuilds a parent as soon as a child changes, so the assertion fired before any rule could see the parent and aborted the replan (apache#2047). NormalizeInterleaveRule rewrites interleaves to UnionExec at the head of each replan; EnforceDistribution re-forms them later through its can_interleave-guarded constructor. A replan can drop a stage that still has tasks in flight. Their completions hit an error that the task-status handler only logs, so the stage never advanced and the job hung forever. Late completions for superseded stages are now ignored. DataFusion's join_selection promotes a small build side to CollectLeft without restricting by join type. That is safe in one process, but Ballista runs one task per probe partition, each with a full copy of the build side, so a Left join emitted its unmatched rows once per task. TPC-DS q77 summed a 16x inflated column as a result. The static planner already guards this in maybe_promote_to_broadcast; DemoteUnsafeBroadcastJoinRule applies the same guard to AQE, between join_selection and EnsureRequirements. ExchangeExec inherited its input's equivalence properties, including constants that only hold within a partition. A UNION ALL branch projecting a literal marks that column constant, and keeping it across a hash repartition on the same column let EnforceSorting drop it from a sort whose SortPreservingMergeExec still merged on it, so q76 returned the wrong rows. RepartitionExec clears these for the same reason. Finally, LateCollectLeft broadcast a build side after both inputs had already been shuffled, paying for both. It is now restricted to null-aware anti joins, which require single-task CollectLeft. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
q31 and q71 were skipped because ORDER BY leaves the order of rows with equal sort keys unspecified, so two correct engines can emit tied rows in different orders and a positional diff is unstable. Compare positionally first, which keeps ordering under test, and only fall back to a canonical row-sorted comparison once that has failed. A wrong value still fails; a pure permutation passes and is reported as such. The skip list is now empty. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
AQE chooses joins and partition counts from measured runtime statistics rather than planning-time estimates. With the correctness fixes in place it verifies clean against single-process DataFusion across all 99 TPC-DS queries, and cuts the TPC-H SF10 suite from 62.2s to 29.6s (-52%). Physical-plan submissions now fall back to the static planner instead of erroring, since AQE plans from the logical plan and that path would otherwise break on the default configuration. Adds SessionConfigExt::with_ballista_adaptive_query_planner so the static planner can be selected in code. Scheduler tests that assert the static planner's stage and partition layout pin it explicitly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Adding one more fix on top to solve the AQE regression. |
|
Nice, the regression fix is a -20% win (and makes every TPC-H query faster under AQE): |
AQE wrapped both inputs of every join in a fresh exchange before resolving it, because repartitioning is what gives the resolver measured statistics. An equi-join already leaves its output partitioned on both sides of each join key, so a downstream join on the other key is co-partitioned and has nothing to wait for. TPC-H q18 moved lineitem twice as a result, 166.5M rows shuffled against the static planner's 106.5M, and shuffle-write CPU doubled from 12.3s to 24.4s. When both inputs already satisfy the required distribution, treat the join as repartitioned and resolve it in the same pass. Partitioning::satisfaction consults equivalence classes, so it sees that a join on l_orderkey = o_orderkey leaves the output partitioned on both. Skipping the exchange without resolving in the same pass does not work: the exchange is also the stage boundary that keeps an unresolved DynamicJoinSelectionExec out of a dispatched stage, and the stage plan then fails to serialize. TPC-H SF10 drops from 29.6s to 23.8s (-20%), q18 from 8.0s to 3.5s (-56%), and TPC-DS still verifies 99/99 against single-process DataFusion. The EXPLAIN tests assert the static planner's stage layout, so they now pin it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
| #[case] | ||
| ctx: SessionContext, | ||
| ) { | ||
| // Golden plan text is the static planner's stage layout. |
There was a problem hiding this comment.
note for follow up: we need to check EXPLAN format with AQE
| estimates. Set to false to use the static distributed planner.".to_string(), | ||
| DataType::Boolean, | ||
| Some(false.to_string())), | ||
| Some(true.to_string())), |
There was a problem hiding this comment.
do we want to do TPCH100 or TPCH1K before we enable AQE?
There was a problem hiding this comment.
yeah asked @andygrove if he would like to run it with higher scale-factors (I tried only SF=10 locally).
There was a problem hiding this comment.
I'll get this done today, hopefully, or at least in the next couple days
|
|
||
| let config = | ||
| SessionConfig::new_with_ballista().with_target_partitions(total_vcores); | ||
| let config = SessionConfig::new_with_ballista() |
There was a problem hiding this comment.
yes this one was failing as empty rule removes things
…default Two configuration defaults change, both aimed at the per-stage overhead that dominates short distributed queries. `ballista.scheduler.max_partitions_per_task` moves from `1` to `0` (unbounded). At `1` every input partition became its own task, so a 16-partition stage dispatched 16 tasks, each carrying its own restricted plan through protobuf encode, dispatch and decode. At `0` the scheduler fills each task up to the assigned executor's free vcore count instead, and operators such as sort and hash join can work across partitions inside one plan invocation. `ballista.optimizer.broadcast_join_threshold_bytes` moves from 10 MB to 128 MB, promoting more joins to a broadcast `CollectLeft` join. A broadcast join removes a shuffle outright: no hash-repartition, no shuffle files written or read back. The cost is that the build side is replicated into every concurrent probe task, so peak memory per executor scales with the number of tasks it runs at once. Measured on TPC-H SF10 (2 executors x 4 vcores, target_partitions=16), running each configuration once per round with the order rotated between rounds so drift is not attributed to a config. Against that baseline the two settings together were the best combination, and they compose — they attack different costs. Enabling the coalesce rule or forcing hash joins on top added little. Both defaults are documented in the 55.0.0 upgrade guide along with how to restore the previous behavior. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
milenkovicm
left a comment
There was a problem hiding this comment.
thanks @Dandandan for addressing AQE issues, great to see that AQE work starts to pay off.
If you're happy with AQE to be on by default i would not object, but lets wait for @andygrove and @avantgardnerio opinion
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
I've run this PR with AQE enabled (local MBP) cargo run --bin tpch -- benchmark ballista -p /Users/marko/TMP/tpch_data/tpch-data-sf10/ -f parquet -i 1 --port 50050 --host 127.0.0.1 -c datafusion.execution.target_partitions=14And with AQE disabled cargo run --bin tpch -- benchmark ballista -p /Users/marko/TMP/tpch_data/tpch-data-sf10/ -f parquet -i 1 --port 50050 --host 127.0.0.1 -c datafusion.execution.target_partitions=14 -c ballista.planner.adaptive.enabled=falseI remember when total time was like ~45 seconds, huge improvement in last few versions, very happy to see this |
I'll wait some days for @andygrove and @avantgardnerio I also found some easy ~17% improvement by enabling:
|
|
I'll run tpc-h 1000 on it tonight or tomorrow. From what I last saw, the perf regression suite is broken in main, and Andy's last test showed MPTs were a regression at that scale. I think it warrants another try with these fixes though. |
`perf(core): pack partitions into tasks and broadcast up to 128 MB by default` raised `ballista.optimizer.broadcast_join_threshold_bytes` from 10 MB to 128 MB but left the tests that read the shipped value behind. * `context_checks`: `should_set_collect_left_thresholds` asserted the old 10485760 for `datafusion.optimizer.hash_join_single_partition_threshold`. This is the failure CI reported. * `broadcast_thresholds`: the 64 MB "known size over threshold" case now sits *under* the threshold, and `wide_schema` (24 bytes/row) can no longer produce a build side that is under the 1,000,000-row threshold yet over 128 MB. Widen the schema to two `Binary` payload columns (204 bytes/row, so 800k rows is ~163 MB) and raise the known-size case to 256 MB. Both queries now read the payload columns, since projection pushdown trims the scan the estimate is taken from. * The two null-aware anti join tests used a 20 MB build side to prove an oversized side is rejected; raised to 256 MB. Also corrects comments and the 55.0.0 upgrade guide, which still described the default as 10 MB and "unchanged". Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`./dev/update_config_docs.sh --check` failed: the generated table in configs.md was hand-edited when the broadcast default changed, leaving the description column padded to a width the generator no longer emits. Regenerated; the change is column padding only. The hand-maintained AQE table in the tuning guide was also left behind by `feat(core): enable adaptive query planning by default`, which updated only configs.md. It still documented `ballista.planner.adaptive.enabled` as defaulting to `false`, and described a job-submission warning that no longer exists in the scheduler. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Where do you see this? I looked at recent commits to main and see the tests passing |
Drops the `SET ballista.planner.adaptive.enabled = false` pins from the EXPLAIN and EXPLAIN ANALYZE tests so they exercise the planner that is now the default, and updates the golden text to the adaptive layout: no `RepartitionExec` above the partial aggregate, and stages numbered from 0 rather than 1. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Tried this PR on our EKS cluster running TPC-H SF1000 (parquet/zstd, 32 executors × 8 concurrent tasks, 2 iterations per query). AQE default-on as this PR flips it. Total wall time: 632.7s vs 721.7s on our most recent baseline before this PR (same cluster, same config) — ~12% faster overall. Q11 now passes (15.2s). It was erroring on every prior run for us, so the AQE plan-rewrite fix here appears to resolve it directly. Per-query deltas vs baseline (bold = >10% change):
New failure — Q10: Q21 and Q22 also errored, but they fail on every recent run for us with connection/h2 errors unrelated to this PR — likely a cluster-side scheduler issue at end-of-run. Not counting them against this PR. Net: strong overall win, one query fixed, one query regressed with cause still to be identified. I'll follow up with the Q10 trace. |
andygrove
left a comment
There was a problem hiding this comment.
LGTM. Thanks @Dandandan! This is awesome
|
Hmm a few tests seems to be hanging, I'll find out why |
|
Ah found it I think |
A task that packs several input partitions drives them all through one plan instance, so anything that plan builds once and shares between partitions -- a CollectLeft join's build side, held in DataFusion's OnceFut -- is polled by whichever write reaches it first on behalf of every other write. futures' Shared only wakes its waiters on the Poll::Ready path. When a poll panics, Reset::drop marks the state POISONED and returns without waking anyone, so every write already parked on that build side stays parked forever. Both shuffle writers joined their per-partition tasks in partition order, so the coordinator blocked on a sibling that could never finish: the task never reported a status, the scheduler kept waiting on it, and the job hung instead of failing. Join the writes as they complete and abandon the rest on the first failure. The task's output is discarded once any write has failed, so there is nothing left to wait for, and dropping the JoinSet aborts whatever is still parked. Reachable since max_partitions_per_task moved to 0: at one partition per task each write had its own plan instance and its own OnceFut, so a panic could only affect the write that raised it. The chaos scenario panicking_task_fails_the_job_but_the_executor_survives hung under AQE, which plans the join as CollectLeft in the same stage as the panicking filter, and that one wedged test stalled every other test behind the harness's process-wide cluster lock -- and with it the whole CI job. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The reasoning is in the fix commit; the code does not need it repeated. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Follow-up on my earlier report (#2315 (comment)): fixed a truncation bug in our runner and got the real Q10 trace — it's a memory-pool OOM in Filed as: #2321 |
Summary
Fixes five defects in adaptive query planning (AQE), removes a redundant shuffle
it inserted, and enables it by default.
Before this change the TPC-DS suite hung at q23 with AQE on, and two queries
returned wrong data. It now verifies clean against single-process DataFusion
across all 99 queries.
Bugs fixed
InterleaveExecinvariant broken by per-branch rewrites (fixes #2047).InterleaveExecrequires all children to share oneHash/Rangepartitioningand asserts on rebuild. AQE rewrites union branches independently, and
TreeNoderebuilds a parent as soon as a child changes, so the assertion fired before any
rule could inspect the parent and aborted the whole replan. A new
NormalizeInterleaveRulerewrites interleaves toUnionExecat the head of eachreplan;
EnforceDistributionre-forms them later through itscan_interleave-guarded constructor.Stale stage completions hung the job. A replan can drop a stage that still
has tasks in flight. Their completions returned an error that the task-status
handler only logs, so the stage never advanced and the query blocked forever.
Late completions for superseded stages are now ignored.
Broadcast-unsafe
CollectLeftjoins produced duplicate rows. DataFusion'sjoin_selectionpromotes a small build side toCollectLeftwithout restrictingby join type. That is correct in one process, but Ballista runs one task per
probe partition, each with a full copy of the build side, so a
Leftjoinemitted its unmatched rows once per task. TPC-DS q77 reported a 16x inflated
salescolumn as a result. The static planner already guards this inmaybe_promote_to_broadcast;DemoteUnsafeBroadcastJoinRuleapplies the sameguard to AQE, between
join_selectionandEnsureRequirements.ExchangeExecleaked per-partition constants across a shuffle. It inheritedits input's equivalence properties, including constants that only hold within a
partition. A
UNION ALLbranch projecting a literal marks that column constant,and keeping it across a hash repartition on the same column let
EnforceSortingdrop it from a sort whose
SortPreservingMergeExecstill merged on it, so q76returned the wrong rows.
RepartitionExecclears these for the same reason.LateCollectLeftbroadcast after shuffling. It paid for the shuffle and thebroadcast, and re-read the build side into every probe task. It is now
restricted to null-aware anti joins, which require single-task
CollectLeft.Redundant shuffle
AQE wrapped both inputs of every join in a fresh exchange before resolving it,
because repartitioning is what gives the resolver measured statistics. An
equi-join already leaves its output partitioned on both sides of each join key,
so a downstream join on the other key is co-partitioned and has nothing to wait
for. TPC-H q18 shuffled
lineitemtwice as a result, moving 166.5M rows againstthe static planner's 106.5M.
When both inputs already satisfy the required distribution the join is now
resolved in the same pass, with no new exchange.
Partitioning::satisfactionconsults equivalence classes, so it sees that a join on
l_orderkey = o_orderkeyleaves the output partitioned on both. Skipping theexchange without resolving in the same pass does not work: the exchange is also
the stage boundary that keeps an unresolved
DynamicJoinSelectionExecout of adispatched stage.
Verification
benchmarks/src/bin/tpcds.rspreviously skipped q31 and q71 becauseORDER BYleaves the order of rows with equal sort keys unspecified, so a positional diff
is unstable. The comparison now falls back to a canonical row-sorted check only
after the positional one fails, so a wrong value still fails but a pure
permutation passes. The skip list is empty and all 99 queries are verified.
TPC-DS SF1, AQE on, each query diffed against single-process DataFusion:
Performance
TPC-H SF10, one executor, 10 vcores, 16 GB pool,
--partitions 10, best of 3.62% faster overall, and every query except q1 (+0.1 s, within noise) improves.
Largest gains are q8 -91%, q9 -87%, q11 -86%, q2 -76%, q17 -73%, q5 -56%,
q21 -55%.
Removing the redundant shuffle accounts for 20% of that on its own (29.6 s to
23.8 s), with q18 -56%, q22 -25%, q19 -21% and q21 -18%.
The static planner path is unchanged (+1.6%, within noise), and the correctness
fixes on their own are performance neutral (-0.2%).
Task packing and broadcast defaults
Two configuration defaults change alongside the AQE work, both aimed at the
per-stage overhead that dominates short distributed queries.
ballista.scheduler.max_partitions_per_taskmoves from1to0(unbounded).At
1every input partition became its own task, so a 16-partition stagedispatched 16 tasks, each carrying its own restricted plan through protobuf
encode, dispatch and decode. At
0the scheduler fills each task up to theassigned executor's free vcore count instead, and operators such as sort and
hash join can work across partitions inside one plan invocation.
ballista.optimizer.broadcast_join_threshold_bytesmoves from 10 MB to 128 MB.A broadcast join removes a shuffle outright: no hash-repartition, and no shuffle
files written and read back. The cost is that the build side is replicated into
every concurrent probe task, so peak memory per executor scales with the number
of tasks it runs at once. The same value drives DataFusion's
hash_join_single_partition_threshold, so the static planner tracks it too.TPC-H SF10, 2 executors x 4 vcores,
--partitions 16. Each configuration runsonce per round with the order rotated between rounds, so drift is spread across
configurations rather than landing on whichever ran last; the table reports the
median of 3 rounds.
The two compose: they attack different costs and neither subsumes the other.
Two further combinations were measured and are deliberately left out. Enabling
the AQE coalesce rule on top reaches 19.70 s (-17.8%), but it trades downstream
parallelism for fewer, larger tasks and currently bails out on broadcast leaves,
range-repartitioned leaves and heterogeneous partition counts, so it deserves its
own change. Forcing hash joins on top of that
(
hash_join_max_build_partition_bytes = 0) regresses to 20.41 s, so sort-mergeremains the right default.
Both defaults are documented in the 55.0.0 upgrade guide, including the higher
per-task memory use and how to restore the previous behavior.
Notes
Physical-plan submissions now fall back to the static planner rather than
erroring, since AQE plans from the logical plan and that path would otherwise
break on the default configuration.
SessionConfigExt::with_ballista_adaptive_query_plannerselects the static planner in code, and scheduler tests that assert the static
planner's stage layout pin it explicitly.
🤖 Generated with Claude Code