feat: spilling hash join operator (opt-in, inner joins) - #2108
Draft
andygrove wants to merge 15 commits into
Draft
feat: spilling hash join operator (opt-in, inner joins)#2108andygrove wants to merge 15 commits into
andygrove wants to merge 15 commits into
Conversation
Add ProbeTable, the correctness core of the spilling hash join: builds an inner-join hash table over one resident build-side bucket and probes it against a probe-side batch, returning matched (build_row, probe_row) index pairs. Also adds assemble_output to materialize the joined RecordBatch from those index pairs. Key equality is resolved via Arrow's row encoding (RowConverter/Rows) rather than a per-DataType match, so Int32/Int64/Decimal128/Date32/Utf8 (and other row-encodable types) all work without dedicated dispatch code. A shared per-row null-key mask enforces NullEquality::NullEqualsNothing semantics, since Arrow's row encoding otherwise treats two NULLs as equal.
Wire SpillingHashJoinExec::execute() to a real join stream: drain the build (left) side fully into one resident ProbeTable per sub-partition bucket, then probe the right side against those tables as it streams in. All build buckets stay resident for now; spilling any of them to disk under memory pressure is left to a later change. The stream is built as a futures::stream::once future (drain + build) that resolves to the probe stream and is flattened with try_flatten, the same idiom SortShuffleWriterExec::execute already uses in this crate for "do async setup, then hand back a stream" — this keeps the build phase as plain async/await instead of a hand-rolled poll_next state machine. Adds an oracle test comparing SpillingHashJoinExec against DataFusion's own HashJoinExec (Inner, Partitioned) over the same hash-repartitioned inputs, so both joins see keys co-located the same way and the comparison actually exercises the Partitioned-mode contract instead of passing vacuously.
Spill both sides of SpillingHashJoinExec to disk and drain spilled buckets one at a time so the join runs in memory bounded by the runtime MemoryPool. - Add per-bucket Arrow-IPC spill files (JoinSpillWriter/JoinSpillReader), with independent build-side and probe-side spill sets per task. - Build phase reserves resident bucket memory against the pool and, on a rejected try_grow, evicts the largest resident bucket to disk; once a bucket cannot fit it stays spilled. - Probe phase probes resident buckets and routes spilled buckets' probe rows to disk; the drain phase reads each spilled bucket back one at a time, recomputing hashes from the fixed partitioner seed, and drops each table before the next so peak drain memory is one bucket. - Publish spill_count/spilled_bytes metrics; gate with tiny-pool oracle tests asserting output identical to a generous-pool HashJoinExec.
Pool-track the drain-phase build side of SpillingHashJoinExec and, when a single spilled bucket's build side will not fit in memory, recursively re-partition both its build and probe spills into finer sub-buckets under a depth-varied hash seed (build and probe share the seed at each level so equal keys stay co-located). Only one (sub-)bucket's build table is resident at any depth. A build side for a single join key that exceeds the whole pool is reported as a clean error instead of an OOM; a best-effort infallible reservation covers residuals that fit the pool but lost a grow race with other resident buckets. A depth cap bounds the recursion.
Add a SpillingHashJoinExecNode proto message and codec arms so SpillingHashJoinExec survives the scheduler->executor plan round trip. Following the ChaosExec convention, children are not embedded in the proto message; datafusion-proto decodes them and passes them through try_decode's inputs parameter. Join keys are serialized via the same physical-expr helpers already used for shuffle hash-repartition exprs.
…dtrip The roundtrip test previously built left and right inputs from the same schema with identical Column exprs on both sides of each on pair, so a bug swapping left_keys/right_keys through encode/decode would still produce byte-identical Display output. Give left and right distinct schemas (l_* vs r_* column names) so a side swap or mis-pairing now changes the rendered on=[...] string and fails the test.
…e joins Add SpillingHashJoinRule, a PhysicalOptimizerRule that swaps eligible HashJoinExec nodes (Inner join, Partitioned mode, no projection, no residual filter) for SpillingHashJoinExec. The rule is gated by ballista.execution.spilling_hash_join.enabled read from the BallistaConfig extension on ConfigOptions, defaulting to disabled when the extension is absent, and uses ballista.execution.spilling_hash_join.partitions (default 16) for the number of spill sub-partitions.
…ast promotion Broadcast promotion only matches HashJoinExec, so a join substituted to SpillingHashJoinExec up front could never be promoted to a shuffle-avoiding CollectLeft broadcast. Move the substitution to run per node immediately after broadcast promotion inside plan_query_stages_internal, so a small build side is offered a broadcast first and only joins left as Inner+Partitioned HashJoinExec are substituted. Factor the single-node eligibility and substitution into a shared maybe_substitute_spilling_hash_join function reused by both SpillingHashJoinRule and the distributed planner, keeping one definition of eligibility.
…lic api Close four correctness/API gaps found in whole-branch review: - Reject substitution when the source HashJoinExec uses NullEquality::NullEqualsNull (e.g. IS NOT DISTINCT FROM); the replacement operator hard-implements NullEqualsNothing semantics, so substituting here would silently drop null-key matches with no error. - Reject non-Partitioned partition_mode in try_new; execute() always runs Partitioned semantics, so a constructed CollectLeft instance would silently mis-execute. - Reject num_sub_partitions == 0 in try_new, before it can reach the build-side hash bucketing and panic on division by zero. - Narrow ProbeTable/assemble_output/RowPartitioner/PartitionedBatch out of the crate's public API; nothing outside the module reaches them through the mod-level re-export, only stream.rs's existing super:: paths.
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
1 participant
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?
Relates to #2025.
This does not
Closes#2025: the operator introduced here is opt-in and disabledby default, so it does not change the default behaviour that OOMs. It provides a
mechanism that resolves that OOM when enabled, and is intended as the first step
toward addressing the issue.
Rationale for this change
DataFusion's
HashJoinExecdoes not spill. Its build side is collected fullyinto memory; when a
Partitionedbuild partition exceeds the memory reservationit returns
ResourcesExhausted(or the process OOMs). This is the failure in#2025 — a
Partitionedbuild side on TPC-H Q18 that cannot fit and cannot spill,and does not shrink with
target_partitions.This PR adds a from-scratch spilling hybrid hash join operator for Ballista.
It runs as a plain in-memory hash join when the build side fits, and degrades
gracefully to disk when it does not, so a large build side completes instead of
failing. It is a hybrid join (as in Spark/Postgres), not pure Grace: buckets
that fit stay resident and only the overflow is spilled, so the in-memory fast
path pays no spill cost.
What changes are included in this PR?
A new
SpillingHashJoinExecoperator (ballista-core) plus the schedulerplumbing to substitute it for eligible joins.
Operator (
ballista/core/src/execution_plans/spilling_hash_join/):Pbuckets by a hash of the joinkeys (fixed seed, independent of the shuffle hash).
MemoryReservation; whentry_growfails, the largest resident bucket is spilled to disk (Arrow IPC viathe runtime
DiskManager). Resident buckets are probed streaming.probe stream ends, spilled buckets are drained one at a time (build table
built, probed, dropped) so peak memory is bounded to a single bucket.
re-partitioned with a depth-varied seed; an irreducible single-join-key bucket
returns a clear error rather than OOMing, panicking, or hanging.
BallistaPhysicalExtensionCodec.Scheduler:
SpillingHashJoinRule/maybe_substitute_spilling_hash_joinsubstitutes aneligible
HashJoinExecfor the spilling operator. Eligibility is strict:Innerjoin,PartitionMode::Partitioned, no projection, no residual filter,and
NullEquality::NullEqualsNothing.DefaultDistributedPlanner) and AQE planners,after broadcast promotion — so a small build side that can be promoted to a
CollectLeftbroadcast keeps that shuffle-avoiding plan (aCollectLeftjoinis not
Partitioned, hence ineligible), and only the remaining largePartitionedhash joins are made spillable.Config:
ballista.execution.spilling_hash_join.enabled(defaultfalse) andballista.execution.spilling_hash_join.partitions(default16).Scope (v1): inner joins only. All other join types are left to the existing
planning path unchanged. Outer/semi/anti and broadcast-mode support are follow-ups.
Testing: the operator is validated against a DataFusion
HashJoinExec(
Inner,Partitioned) oracle on hash-co-located inputs, with output comparedunder (a) fully resident, (b) forced build-side spill, (c) forced two-sided
spill, and (d) forced recursive re-partition — down to a ~1 KiB memory pool — plus
a clean-error test for single-key skew. Codec round-trip and planner-substitution
(including broadcast-wins and AQE dynamic-join-not-clobbered cases) are covered by
plan-string tests. TPC-H plan-stability goldens are unchanged (the feature is a
strict no-op when disabled).
Are there any user-facing changes?
Two new configuration settings (above), both defaulting to off/existing behaviour.
No public API changes and no change to default planning. With the feature disabled
(the default), behaviour is unchanged.