jayzhan211 opened a new pull request, #25250: URL: https://github.com/apache/datafusion/pull/25250
## 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`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
