jayzhan211 opened a new pull request, #25986: URL: https://github.com/apache/datafusion/pull/25986
## Which issue does this PR close? - No separate issue. Follow-up to #25206, which did this for the materializing sort-merge join stream (Inner/Left/Right/Full). ## Rationale for this change A semi/anti/mark sort-merge join keeps memory reserved after it can no longer use it, until the stream is dropped. That covers its final output, and for anti/mark joins also every outer batch drained after the inner side runs out. A parent operator in the same query that needs that memory while consuming the join's output can then fail to allocate it under a tight pool. Three cases, all reproduced on `main`: - With a join filter, the reservation for the last buffered inner key group is held. Example: outer `k = [1, 2]`, inner has 32 x 4096 rows with `k = 2`, and a filter that rejects them all. Under a 2 MB pool with spilling disabled, the pool still reports 1,582,080 bytes reserved while the final batch is in hand. A parent `Inner` sort-merge join over 10 x 4096 rows then fails with `Failed to allocate additional 96.4 KB for SMJStream[0] ... Disk spilling disabled` for LeftSemi/LeftAnti/LeftMark, whichever input runs out first. - When the outer side runs out first, the inner input is not dropped, so a `SortExec` below it keeps its unread sorted batches reserved (98,304 bytes in the test). - When the inner side runs out first, a semi join stops reading the outer side (#25206) but does not drop it, so a `SortExec` below the outer side keeps its unread batches reserved in the same way. ## What changes are included in this PR? In `bitwise_stream.rs`: - New `release_inner()`: clears the inner key group, frees the reservation, and replaces the inner input with an `EmptyRecordBatchStream`. - `drain_outer` calls it before emitting the remaining outer batches. A semi join also replaces the outer input with an `EmptyRecordBatchStream`, since it reads no more of it. - `join()` calls it after the merge loop, before the final coalescer flush, which covers the outer side running out first. Between key groups, the reservation is still not shrunk (#20729). ## What is the testing strategy for this PR? New tests in `joins/sort_merge_join/tests.rs`, next to the #25206 tests. All three fail on `main`: - `bitwise_join_releases_inner_key_group_reservation_before_output`: LeftSemi/LeftAnti/LeftMark with the filter above, with either the outer or the inner side running out first. Asserts the pool holds no reservation while any output batch is in hand. On `main` it holds 1,582,080 bytes. - `bitwise_join_chain_reuses_memory_released_by_completed_child`: the same joins below a parent `Inner` join, 2 MB pool, spilling disabled. On `main` the parent fails to allocate its buffered key group. - `bitwise_join_releases_unread_input_before_final_batch`: puts a `SortExec` under the outer side (inner runs out first) or under the inner side (outer runs out first), and asserts nothing is reserved at the final batch. On `main`, 98,304 bytes are still reserved for LeftSemi in the first shape and for every join type in the second. All existing `sort_merge_join` tests, the join fuzz tests and the sqllogictests pass. #25933 also changes `bitwise_stream.rs` (one line in `buffer_inner_key_group`) and adds tests to `tests.rs` in a different place. The two PRs do not conflict. ## Are there any user-facing changes? No API or result changes. Semi/anti/mark sort-merge joins give memory back to the pool earlier. -- 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]
