Skip to content

fix: register the sort-merge join stream as a spillable memory consumer - #25250

Merged
jayzhan211 merged 2 commits into
apache:mainfrom
jayzhan211:fix/smj-stream-can-spill
Sep 16, 2026
Merged

jayzhan211 merged 2 commits into
apache:mainfrom
jayzhan211:fix/smj-stream-can-spill

Conversation

@jayzhan211

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

  • No existing issue. Found while preparing the hash join sort-merge fallback, which runs this stream next to two external sorts inside one partition.

Rationale for this change

SMJStream spills: when try_grow for a buffered batch fails it writes the buffered side to disk through its SpillManager and carries on. But it registers its memory consumer with the default can_spill = false, so the memory pool is told the opposite.

Only FairSpillPool reads that flag, and it uses it to run two different budgets:

  • A consumer that can spill is limited to an even share, (pool_size - unspillable) / num_spillable, and is expected to spill when it reaches it.
  • A consumer that cannot spill takes free memory first come, first served, and everything it holds is subtracted from the pool before the spillable shares are computed.

Registering the join as unspillable therefore has three effects, all wrong for an operator that spills:

  1. It is not counted in num_spillable, so the sorts feeding it split the pool as if the join needed nothing.
  2. Every buffered batch it holds shrinks the sorts' shares, since it is booked as unspillable memory.
  3. It is never asked to spill until the whole pool is allocated.

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. FairSpillPool users are affected, which includes datafusion-cli --mem-pool-type fair.

What changes are included in this PR?

One line in SortMergeJoinExec::execute: .with_can_spill(true) on the stream's MemoryConsumer, 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_consumer in sort_merge_join/tests.rs runs a small join against a recording pool that notes each consumer's name and can_spill flag 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 default GreedyMemoryPool.

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 13, 2026
@codecov-commenter

codecov-commenter commented Sep 13, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.91%. Comparing base (a407990) to head (7a4cbf1).
⚠️ Report is 38 commits behind head on main.

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@jayzhan211
jayzhan211 marked this pull request as ready for review September 13, 2026 05:07
// 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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The registration dates from #5632, when the join could not spill and the flag was accurate; #11218 added spilling without updating it.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

  1. The join is not counted as a spillable consumer, so the sorts feeding it split the pool as if the join needed nothing.
  2. Every batch the join buffers is booked as unspillable, so it shrinks every other spillable operator's share as it grows.
  3. 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.

@jayzhan211
jayzhan211 requested a review from kosiew September 14, 2026 13:42

@kosiew kosiew left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jayzhan211,

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<()> {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@jayzhan211
jayzhan211 added this pull request to the merge queue Sep 16, 2026
Merged via the queue into apache:main with commit 767d142 Sep 16, 2026
41 checks passed
@jayzhan211
jayzhan211 deleted the fix/smj-stream-can-spill branch September 16, 2026 00:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants