fix: register the sort-merge join stream as a spillable memory consumer - #25250
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #25250 +/- ##
==========================================
- Coverage 81.95% 81.91% -0.05%
==========================================
Files 1133 1135 +2
Lines 423828 427417 +3589
Branches 423828 427417 +3589
==========================================
+ Hits 347344 350113 +2769
- Misses 55890 56378 +488
- Partials 20594 20926 +332 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| // pool that budgets spillable and unspillable consumers differently | ||
| // (`FairSpillPool`) has to know it can. | ||
| let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]")) | ||
| .with_can_spill(true) |
There was a problem hiding this comment.
What if we keep it false
Under GreedyMemoryPool or UnboundedMemoryPool: nothing. Neither reads the flag. A plain memory limit gives you the greedy pool, so most users would never notice.
Under FairSpillPool, three things go wrong, all from one cause. The pool budgets a consumer that cannot spill differently from one that can: it lets it take free memory first come, first served, and subtracts whatever it holds before dividing the rest evenly among the spillable consumers.
- The join is not counted as a spillable consumer, so the sorts feeding it split the pool as if the join needed nothing.
- Every batch the join buffers is booked as unspillable, so it shrinks every other spillable operator's share as it grows.
- The join is never asked to spill until the whole pool is allocated.
What that looks like in practice. Take a 300 MB pool with two sorts feeding the join. Registered as spillable, each of the three gets 100 MB. Registered as unspillable, the sorts each get half of whatever the join has not taken. If the join buffers 280 MB of one large key group, the sorts are down to 10 MB each, which is the default sort_spill_reservation_bytes a sort must hold to merge its spill files. Below that, the sort fails with Resources exhausted, and the query fails inside an operator that was behaving correctly, when the join could have spilled instead. Short of failure, the sorts spill more than necessary.
It also makes the outcome depend on timing. If the sorts hold memory first, the join sees only the leftovers and spills early and often. If the join buffers first, the sorts starve. Removing exactly that order dependence is what FairSpillPool exists for.
For the fallback specifically, this is the shape every fallen-back partition has: two external sorts plus this stream, sharing one pool. Keeping it false would make the fallback's memory behavior under the fair pool unfair in the same way.
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this. The change looks good to me. I left one optional suggestion for strengthening the regression coverage around the FairSpillPool behavior.
| /// split, and may take everything they have not yet claimed, starving the | ||
| /// sorts the join usually runs on top of. | ||
| #[tokio::test] | ||
| async fn stream_registers_as_a_spillable_consumer() -> Result<()> { |
There was a problem hiding this comment.
Nice, this recording pool directly verifies the registration metadata changed by this PR. One optional improvement would be to add a stronger regression test using a constrained FairSpillPool, with the SMJ and at least one other concurrently registered spillable consumer, such as the sorts feeding it. The test could then verify that the join completes and that the SMJ actually records a spill.
A large equal-key group by itself would not quite cover the regression, since the SMJ can spill and complete both before and after this change. Ideally, the test should fail without .with_can_spill(true) so it exercises the fair-share behavior this fix is addressing.
Which issue does this PR close?
Rationale for this change
SMJStreamspills: whentry_growfor a buffered batch fails it writes the buffered side to disk through itsSpillManagerand carries on. But it registers its memory consumer with the defaultcan_spill = false, so the memory pool is told the opposite.Only
FairSpillPoolreads that flag, and it uses it to run two different budgets:(pool_size - unspillable) / num_spillable, and is expected to spill when it reaches it.Registering the join as unspillable therefore has three effects, all wrong for an operator that spills:
num_spillable, so the sorts feeding it split the pool as if the join needed nothing.So the operator that could spill cheaply squeezes the operators beside it. As the join buffers a large key group, the sorts' shares fall toward zero, they spill more than needed, and once a share drops below what a sort must hold for its merge (
sort_spill_reservation_bytes) the sort fails outright. A query that would have finished by spilling the join instead fails in the sort.With the flag set, the join takes an even share, its buffered memory is booked as spillable, and it spills when it exceeds that share, which is the contract the pool's documentation describes for spillable operators.
Blast radius:
GreedyMemoryPool, the default whenever a memory limit is set, ignores the flag, so nothing changes there.FairSpillPoolusers are affected, which includesdatafusion-cli --mem-pool-type fair.What changes are included in this PR?
One line in
SortMergeJoinExec::execute:.with_can_spill(true)on the stream'sMemoryConsumer, with a comment saying why. The bitwise stream used by semi, anti and mark joins registers separately and is not touched here; it may have the same omission and I have not checked.What is the testing strategy for this PR?
stream_registers_as_a_spillable_consumerinsort_merge_join/tests.rsruns a small join against a recording pool that notes each consumer's name andcan_spillflag at registration, then asserts the stream's entry is spillable. With the fix reverted, the test fails on "the sort-merge join stream must register as able to spill".The test deliberately does not run under
FairSpillPool: what a stream may hold there depends on which other consumers are alive at each allocation, so an end-to-end assertion would depend on scheduling. Reading the flag at registration is the deterministic form of the same claim.The existing sort-merge join suite passes unchanged (234 tests), and clippy is clean with all targets and features.
Are there any user-facing changes?
No API change. Under
FairSpillPool, a sort-merge join now spills at its fair share instead of consuming free memory first, so the sorts feeding it keep their budget. No change under the defaultGreedyMemoryPool.