fwojciec opened a new pull request, #26994: URL: https://github.com/apache/datafusion/pull/26994
## Which issue does this PR close? - Closes #26980. ## Rationale for this change A `Partitioned` hash join frees its build reservation when the partition's stream ends, but the dynamic filter and the build accumulator keep the build map until the plan is dropped, so a finished join's maps stay allocated while the pool counts them as free. In a query with several stages, later operators reserve the whole pool on top of them, and the process can be killed by the OS instead of failing with `ResourcesExhausted`. #26980 has the cause and the measurements. This takes the first fix direction from the issue: the reservation travels with the map, so the memory stays charged for as long as anything holds the map. This does not reduce memory use or free the maps sooner. A memory-limited query that only fit because a finished join's maps were not counted can now fail with `ResourcesExhausted`, and the finished join then shows up among the top consumers in the error. ## What changes are included in this PR? - `ReservedMap` (`joins/mod.rs`): a `Map` and the reservation its memory is charged to, in one `Arc`. `JoinBuildData`, `PushdownStrategy::Map` and `HashTableLookupExpr` hold it in place of `Arc<Map>`, so every holder of the map holds its charge. - `build_in_memory`: charges the hash table or array map and the pruning bitmap to a reservation created with `new_empty()` from the build reservation, so the consumer name and error text do not change. The batch, key arrays and visited bitmap stay on the build reservation and are released with the stream as before. - `HashTableLookupExpr::new`: keeps its signature and wraps the map without a reservation, for the `datafusion-proto` and `datafusion-pruning` tests. `new_reserved` is the crate-internal constructor the join uses. - `PreparedHashJoinBuild::reserved_bytes`: sums both reservations, so its value is unchanged. Not in this PR: - Releasing the maps once every partition has finished probing, the second fix direction in #26980, frees the memory sooner but changes when the filter is updated, so it is left for a follow-up. I can open an issue for it. - `PushdownStrategy::InList` arrays share the prepared build batch and have the same lifetime gap, bounded by `hash_join_inlist_pushdown_max_size` (128 KiB per partition by default). ## What is the testing strategy for this PR? - `build_map_stays_reserved_with_dynamic_filter` in `hash_join/exec.rs`, over `CollectLeft` and `Partitioned`. After the join's stream is dropped, `pool.reserved()` equals the map plus pruning bitmap bytes (`Partitioned`) or the whole build (`CollectLeft`). After the plan is dropped, the filter is the last holder and both modes equal the map plus bitmap bytes. After the filter is dropped, it is 0. On `main` the `Partitioned` case fails at the first assertion with 0 reserved against 136,288 expected; `CollectLeft` passes on `main` too. - The test from #26980, run on this branch (not added here: it needs a global allocator): | Run | Allocated after the stream is dropped | Reserved | Reserved on `main` | |---|---|---|---| | 1. `Partitioned`, dynamic filter off | 0.2 MiB | 0.0 MiB | 0.0 MiB | | 2. `Partitioned`, dynamic filter on | 151.5 MiB | 151.3 MiB | 0.0 MiB | | 3. `CollectLeft`, dynamic filter on | 182.0 MiB | 181.8 MiB | 181.8 MiB | Dropping the plan frees everything in all three runs. ## Are there any user-facing changes? No API change; `HashTableLookupExpr::new` is unchanged. Under a memory limit, a query whose finished joins' maps were not counted before can now fail with `ResourcesExhausted`, and spilling operators under the same limit can spill earlier. AI assistance: written with Claude Code (Claude Fable 5.1). I reported the original case, reproduced the issue, reviewed the diff and ran the tests above. -- 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]
