feat(scheduler): adopt DataFusion's cost-based join order enumeration - #2352
Draft
Dandandan wants to merge 3 commits into
Draft
feat(scheduler): adopt DataFusion's cost-based join order enumeration#2352Dandandan wants to merge 3 commits into
Dandandan wants to merge 3 commits into
Conversation
Moves the crates.io patch off the 55.0.0-rc3 tag and onto apache/datafusion#24456, which adds a `JoinEnumeration` physical optimizer rule that searches join orders from cardinality estimates and considers bushy shapes as well as left-deep ones. The pin is a personal fork rev because the PR is not merged yet; repoint it at apache/datafusion once it is. Its merge base with main is f1f0449a53. Two adaptations to DataFusion main, both independent of that PR: * `TableProvider::scan` now takes `Option<&[usize]>` rather than `Option<&Vec<usize>>`. Updates the three test providers. * apache/datafusion#24357 rejects a `RANGE` offset frame over a `Utf8` ORDER BY at type coercion, so the test that checks `ParallelWindowRule` declines a key the sketch cannot encode failed at planning instead of reaching the rule. `SortKeyCodec` covers the integer, float, date, time, timestamp and duration types and nothing else, so order by a `Decimal128` instead: still unencodable, but it accepts a numeric offset frame.
`plan_preparation_optimizers` runs `DelayJoinSelectionRule`, which rewrites every join into a `DynamicJoinSelectionExec` that `SelectJoinRule` unwraps only once its `upstream_resolved()`. DataFusion's default rule list runs after that, both once at plan time and again after every stage completion, so `JoinEnumeration` never sees the whole join graph -- it re-searches a progressively-resolving fragment of it and re-decides each time. On TPC-H SF=10 that cost 45.6% overall and failed q5 outright. With only part of the graph visible the search split q5's composite key `(s_suppkey = l_suppkey, s_nationkey = c_nationkey)` and scheduled `s_nationkey = c_nationkey` as a standalone join. `nationkey` has 25 distinct values, so that join produces 1.83 billion rows against an estimate of 2.3M; the estimate fits under `broadcast_join_threshold_bytes`, so AQE broadcast it and the executor exhausted the whole 22.4 GB pool, 3 runs out of 3. Even join-free q6 regressed 189%, so part of the cost was pure per-replan overhead. Join order is a plan-time decision over the whole graph, so make it once, before the joins are hidden: add the rule to `plan_preparation_optimizers` after `FilterPushdown`, so the search costs each input with its filters already on the scan, and before `DelayJoinSelectionRule`; and exclude it from the replan chain, as `FilterPushdown` already is (apache#2344). Join *implementation* stays deferred to AQE, which is the part that genuinely benefits from measured statistics. TPC-H SF=10 goes from +45.6% to -1.0% and q5 passes at parity.
Two expected effects of the DataFusion bump, neither a behaviour change for Ballista: * apache/datafusion#24456 also raises the default `datafusion.optimizer.hash_join_single_partition_threshold` from 1 MiB to 4 MiB. `alter_stages`' mock scans derive their byte statistics from that default, so 17 recorded sizes scale by 4. Ballista overrides the option with `ballista.optimizer.broadcast_join_threshold_bytes`, so nothing changes at run time. * `hash_join_three_tables_collect_left` now enumerates a different order: the nested join moves to the probe side and a plain scan becomes the build side, so the top projection shifts from `id@0` to `id@1`. Both joins stay `CollectLeft`, which is what the test is about.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
None filed — happy to open one if the project would like this tracked. It adopts
apache/datafusion#24456.
Draft: the
[patch.crates-io]pin points at an unmerged fork rev(
Dandandan/arrow-datafusion@29c224d86b). It needs repointing at anapache/datafusion rev before this can merge.
Rationale for this change
apache/datafusion#24456 adds a
JoinEnumerationphysical optimizer rule thatsearches join orders from cardinality estimates, considering bushy shapes as
well as left-deep ones. Taken as-is it is not usable on Ballista: TPC-H
SF=10 costs 45.6% overall and q5 fails outright, 3 runs out of 3, with
ResourcesExhaustedagainst the whole 22.4 GB executor pool.That is placement, not plan quality.
plan_preparation_optimizersrunsDelayJoinSelectionRule, which rewrites every join into aDynamicJoinSelectionExecthatSelectJoinRuleunwraps only once itsupstream_resolved(). DataFusion's default rule list runs after that — once atplan time (
planner.rs:109) and again after every stage completion(
planner.rs:366) — so the rule never sees the whole join graph. It re-searchesa progressively-resolving fragment and re-decides each time.
With only part of the graph visible the search split q5's composite key
(s_suppkey = l_suppkey, s_nationkey = c_nationkey)and scheduleds_nationkey = c_nationkeyas a standalone join.nationkeyhas 25 distinctvalues, so it produces 1,830,400,706 rows against an estimate of 2.3M / 36.5
MB. That estimate fits under
broadcast_join_threshold_bytes, so AQE broadcastit to every probe task. Plain DataFusion, seeing the whole graph, keeps the key
composite under both
prefer_hash_joinsettings — so this is Ballista-only.Join order is a plan-time decision over the whole graph. Join
implementation is local and genuinely benefits from measured statistics.
Splitting them that way is the fix.
What changes are included in this PR?
chore(deps)— moves the patch off55.0.0-rc3onto the upstream PR,plus two adaptations to DataFusion main that are independent of it:
TableProvider::scannow takesOption<&[usize]>, and fix: Support RANGE window frames over binary ORDER BY keys datafusion#24357rejects a
RANGEoffset frame over aUtf8ORDER BY, so theParallelWindowRuletest that needs an unencodable sort key orders byDecimal128instead.fix(scheduler)— 14 lines inaqe/planner.rs. AddsJoinEnumerationtoplan_preparation_optimizers, afterFilterPushdownso the search costs eachinput with its filters already on the scan, and before
DelayJoinSelectionRuleso it sees real joins; and returnsvec![]for"join_enumeration"indatafusion_optimizers()so it does not re-run perreplan — the same exclusion
FilterPushdownalready has (perf(scheduler): push filters once, not on every replan #2344).test— 17 snapshot byte sizes scale by 4 because the upstream PR alsoraises the default
hash_join_single_partition_thresholdfrom 1 MiB to 4 MiB(inert here, Ballista overrides that option), and
hash_join_three_tables_collect_leftenumerates a different order.Benchmarks
1 scheduler + 2 executors × 4 vcores,
target_partitions=16, 10-core / 32 GBhost. Variant order rotated per round; medians reported. TPC-H is the median of
8 (4 rounds × 2 iterations), TPC-DS the median of 3. Every pass ran to
completion; 0 query failures and 0 row-count mismatches across all 99 TPC-DS
queries.
Baseline is apache/datafusion main at
f1f0449a53, the upstream PR's own mergebase, so the comparison isolates the PR.
Both headline figures have one query doing most of the work, so quote them by
name: TPC-DS q72 goes 32.767 s → 0.248 s (132×), which is 32.5 of the 34.5
seconds saved; TPC-H q18 costs +1.099 s on its own.
TPC-H movers: q7 −55.0%, q3 −51.5%, q12 −35.2%, q2 −26.7%, q16 −13.0%,
q11 −10.7% against q18 +35.0%, q9 +14.7%, q19 +10.6%.
Without commit 2, the same build measures +45.6% on TPC-H SF=10 and fails
q5. Join-free q6 regresses 189% there, so part of that cost is pure per-replan
overhead rather than plan quality.
Known issue: q18
q18 is the one significant regression and it is not a reordering problem —
setting
join_enumeration=falseon this build reproduces the identical 8-stageplan. It is a statistics regression that flips an irreversible AQE decision:
Inexact(1,500,000), bytesAbsentRepartitionExact(15,000,000), bytesExact(871 MB)Hash(Partitioned)Inexact(1,500,000), bytesInexact(99 MB)CollectLeft, never re-measuredBoth estimate the rows equally badly. Baseline had no byte estimate, so the
decision fell to the row path, hit the 1M row ceiling and took
Repartition—which shuffled, measured 871 MB and correctly stayed partitioned. The new 99 MB
estimate undercuts the 128 MiB byte threshold and commits to a broadcast that is
never revisited, which costs the
SinglePartitionedaggregate, theTopKpushdown into that stage, and a 180M-row shuffle on a five-column key.
The root fix is upstream:
customer ⨝ ordersoncustkeyis estimated at 1.5Mrows where it produces 15M, which looks like the PK side's cardinality with no FK
fan-out. Correct it and the byte estimate lands over the threshold, and Ballista
picks partitioned with no change here.
Are there any user-facing changes?
Join plans change, generally for the better. Three new DataFusion options become
available and default on:
datafusion.optimizer.join_enumeration,join_enumeration_min_improvement,join_enumeration_limit. Settingjoin_enumeration=falserestores the previous ordering.No Ballista public API changes.