Skip to content

feat(scheduler): adopt DataFusion's cost-based join order enumeration - #2352

Draft
Dandandan wants to merge 3 commits into
apache:mainfrom
Dandandan:feat/join-order-enumeration
Draft

feat(scheduler): adopt DataFusion's cost-based join order enumeration#2352
Dandandan wants to merge 3 commits into
apache:mainfrom
Dandandan:feat/join-order-enumeration

Conversation

@Dandandan

Copy link
Copy Markdown
Contributor

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 an
apache/datafusion rev before this can merge.

Rationale for this change

apache/datafusion#24456 adds a JoinEnumeration physical optimizer rule that
searches 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
ResourcesExhausted against the whole 22.4 GB executor pool.

That is placement, not plan quality. 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 — once at
plan time (planner.rs:109) and again after every stage completion
(planner.rs:366) — so the rule never sees the whole join graph. It re-searches
a 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 scheduled
s_nationkey = c_nationkey as a standalone join. nationkey has 25 distinct
values, 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 broadcast
it to every probe task. Plain DataFusion, seeing the whole graph, keeps the key
composite under both prefer_hash_join settings — 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?

  1. chore(deps) — moves the patch off 55.0.0-rc3 onto the upstream PR,
    plus two adaptations to DataFusion main that are independent of it:
    TableProvider::scan now takes Option<&[usize]>, and fix: Support RANGE window frames over binary ORDER BY keys datafusion#24357
    rejects a RANGE offset frame over a Utf8 ORDER BY, so the
    ParallelWindowRule test that needs an unencodable sort key orders by
    Decimal128 instead.
  2. fix(scheduler) — 14 lines in aqe/planner.rs. Adds JoinEnumeration to
    plan_preparation_optimizers, after FilterPushdown so the search costs each
    input with its filters already on the scan, and before
    DelayJoinSelectionRule so it sees real joins; and returns vec![] for
    "join_enumeration" in datafusion_optimizers() so it does not re-run per
    replan — the same exclusion FilterPushdown already has (perf(scheduler): push filters once, not on every replan #2344).
  3. test — 17 snapshot byte sizes scale by 4 because the upstream PR also
    raises the default hash_join_single_partition_threshold from 1 MiB to 4 MiB
    (inert here, Ballista overrides that option), and
    hash_join_three_tables_collect_left enumerates a different order.

Benchmarks

1 scheduler + 2 executors × 4 vcores, target_partitions=16, 10-core / 32 GB
host. 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 this PR Δ
TPC-H SF=10, 22 queries 19.623 s 19.423 s −1.0%
… excluding q18 −7.9%
TPC-DS SF=1, 99 queries 58.909 s 24.436 s −58.5%
… excluding q72 26.142 s 24.188 s −7.5%

Baseline is apache/datafusion main at f1f0449a53, the upstream PR's own merge
base, 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=false on this build reproduces the identical 8-stage
plan. It is a statistics regression that flips an irreversible AQE decision:

plan_id=1 build side estimate AQE decision
baseline rows Inexact(1,500,000), bytes Absent rejected → Repartition
→ after measuring rows Exact(15,000,000), bytes Exact(871 MB) Hash(Partitioned)
this PR rows Inexact(1,500,000), bytes Inexact(99 MB) CollectLeft, never re-measured

Both 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 SinglePartitioned aggregate, the TopK
pushdown into that stage, and a 180M-row shuffle on a five-column key.

The root fix is upstream: customer ⨝ orders on custkey is estimated at 1.5M
rows 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. Setting
join_enumeration=false restores the previous ordering.

No Ballista public API changes.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant