jayzhan211 opened a new issue, #26980: URL: https://github.com/apache/datafusion/issues/26980
### Describe the bug A `Partitioned` hash join reserves memory for each partition's hash map and releases that reservation when the partition's stream ends. With `datafusion.optimizer.enable_join_dynamic_filter_pushdown` on (the default), the maps are not freed at that point: the dynamic filter and the join's build accumulator both still hold them, and both live as long as the physical plan. So a finished join's maps stay allocated while the memory pool counts them as free. In a query with several stages, later operators can then reserve the whole pool on top of these maps, and the process can be killed for running out of memory instead of failing with `ResourcesExhausted`. @fwojciec reported this in #24768 from a production workload: about 10 GiB unreserved above a 15–18 GiB pool, then OOM kills. Their standalone repro, which also covers a two-join query: https://github.com/fwojciec/datafusion-join-map-retention-repro ### To Reproduce Save the file below as `datafusion/core/tests/hj_map_retention.rs` and run `cargo test -p datafusion --test hj_map_retention -- --nocapture`. It counts live heap bytes with a global allocator, joins 8M probe rows with 4M build rows over Parquet with 4 partitions, reads the stream to the end, drops the stream, and keeps the plan. <details><summary>hj_map_retention.rs</summary> ```rust //! Memory a hash join leaves allocated after its stream is dropped, next to //! what the memory pool has reserved. Run with: //! cargo test -p datafusion --test hj_map_retention -- --nocapture use std::alloc::{GlobalAlloc, Layout, System}; use std::sync::Arc; use std::sync::atomic::{AtomicIsize, Ordering::Relaxed}; use datafusion::arrow::array::Int64Array; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::parquet::arrow::ArrowWriter; use datafusion::physical_plan::execute_stream; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; use futures::StreamExt; /// Counts live heap bytes. struct Counting; static LIVE: AtomicIsize = AtomicIsize::new(0); unsafe impl GlobalAlloc for Counting { unsafe fn alloc(&self, l: Layout) -> *mut u8 { LIVE.fetch_add(l.size() as isize, Relaxed); unsafe { System.alloc(l) } } unsafe fn dealloc(&self, p: *mut u8, l: Layout) { LIVE.fetch_sub(l.size() as isize, Relaxed); unsafe { System.dealloc(p, l) } } } #[global_allocator] static GLOBAL: Counting = Counting; fn mib(bytes: isize) -> f64 { bytes as f64 / (1024.0 * 1024.0) } /// One Int64 column `k` holding `0, step, 2 * step, ...`, in 4 Parquet files. fn write_table(dir: &std::path::Path, name: &str, rows: i64, step: i64) { let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); std::fs::create_dir_all(dir.join(name)).unwrap(); for f in 0..4 { let path = dir.join(name).join(format!("{f}.parquet")); let file = std::fs::File::create(path).unwrap(); let mut w = ArrowWriter::try_new(file, Arc::clone(&schema), None).unwrap(); let rows = (f * rows / 4..(f + 1) * rows / 4).collect::<Vec<_>>(); for chunk in rows.chunks(65_536) { let col = Int64Array::from_iter_values(chunk.iter().map(|v| v * step)); let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(col)]); w.write(&batch.unwrap()).unwrap(); } w.close().unwrap(); } } async fn run(dir: &std::path::Path, label: &str, repartition: bool, filter: bool) { let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(8 << 30)); let runtime = RuntimeEnvBuilder::new() .with_memory_pool(Arc::clone(&pool)) .build_arc() .unwrap(); let mut config = SessionConfig::new().with_target_partitions(4); let options = config.options_mut(); options.optimizer.enable_join_dynamic_filter_pushdown = filter; // false gives a CollectLeft join. options.optimizer.repartition_joins = repartition; let ctx = SessionContext::new_with_config_rt(config, runtime); for t in ["build", "probe"] { let path = dir.join(t); let options = ParquetReadOptions::default(); ctx.register_parquet(t, path.to_str().unwrap(), options) .await .unwrap(); } let base = LIVE.load(Relaxed); let sql = "SELECT COUNT(*) FROM probe p JOIN build b ON p.k = b.k"; let plan = ctx.sql(sql).await.unwrap().create_physical_plan().await.unwrap(); let mut stream = execute_stream(Arc::clone(&plan), ctx.task_ctx()).unwrap(); while let Some(batch) = stream.next().await { batch.unwrap(); } drop(stream); let allocated = LIVE.load(Relaxed) - base; let reserved = pool.reserved() as isize; drop(plan); let after_plan_drop = LIVE.load(Relaxed) - base; println!( "{label}: after the stream is dropped: allocated {:.1} MiB, reserved {:.1} MiB; \ after the plan is dropped: allocated {:.1} MiB", mib(allocated), mib(reserved), mib(after_plan_drop), ); } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn hj_map_retention() { let dir = tempfile::tempdir().unwrap(); // 4M build keys, every 1000th value, so the join builds a hash table. write_table(dir.path(), "build", 4_000_000, 1000); write_table(dir.path(), "probe", 8_000_000, 1); run(dir.path(), "1. Partitioned, dynamic filter off", true, false).await; run(dir.path(), "2. Partitioned, dynamic filter on ", true, true).await; run(dir.path(), "3. CollectLeft, dynamic filter on ", false, true).await; } ``` </details> ### Expected behavior | Run | Allocated after the stream is dropped | Reserved | Expected | |---|---|---|---| | 1. `Partitioned`, dynamic filter off | 0.2 MiB | 0 MiB | as is | | 2. `Partitioned`, dynamic filter on | 151.5 MiB | 0 MiB | the 151.5 MiB stays reserved, or is freed | | 3. `CollectLeft`, dynamic filter on | 182.0 MiB | 181.8 MiB | as is | Dropping the plan frees everything in all three runs. With dense keys (`step = 1`, which gives an array map instead of a hash table), run 2 keeps 61.2 MiB, close to the 61.3 MiB in fwojciec's repro, which uses similar data. ### Additional context Cause: - The build shares its map with the dynamic filter as `PushdownStrategy::Map(Arc::clone(&map), ...)` ([exec.rs:3368](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/exec.rs#L3368)). The reservation sits in `JoinBuildData` ([exec.rs:209](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/exec.rs#L209)), which a `Partitioned` join drops together with the partition's stream. - The filter holds the map: `create_membership_predicate` wraps it in a `HashTableLookupExpr` ([shared_bounds.rs:128](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs#L128), field `map: Arc<Map>` at [partitioned_hash_eval.rs:272](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs#L272)), inside the dynamic filter, which both the join and the probe-side scan hold. - The accumulator holds the map too: `take_finalize_input_if_ready` clones each partition's data rather than taking it ([shared_bounds.rs:561](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs#L561)), so `SharedBuildAccumulator` keeps every partition's `PushdownStrategy::Map` after the filter is built. The accumulator lives in `HashJoinExecDynamicFilter::build_accumulator` ([exec.rs:881](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/exec.rs#L881)) for as long as the plan does. - `CollectLeft` is not affected: its build data, reservation included, stays in `left_fut` ([exec.rs:814](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/exec.rs#L814)) until the plan is dropped. Both holders matter. In run 2 with local patches, building the filter without the map still leaves 151.5 MiB allocated, and clearing the accumulator's maps after the filter is built also leaves 151.5 MiB. Doing both brings it down to 0.3 MiB. Fix directions: keep a reservation together with the map (for example, one `Arc` that owns both), so the memory stays charged for as long as anything holds the map; this needs no change to the filter. Or drop the map from both holders once every partition has finished probing, which frees the memory sooner. The key range bitmap next to the map is charged to the same reservation and held the same way ([exec.rs:3359](https://github.com/apache/datafusion/blob/5db4b233efde58f6b34299fb15b756e7bdef2698/datafusion/physical-plan/src/joins/hash_join/exec.rs#L3359)). Related: #26442 (draft) replaces the per-partition filter with one global bloom filter, so the filter would no longer hold the maps. As far as I can tell from its diff, the accumulator still keeps them, so by the experiment above that change would not fix this on its own. #25650 tracks OOM kills from inaccurate memory accounting in general. ### Willingness to contribute @fwojciec offered in #24768 to open a PR for this, so they may already be working on it. Otherwise, anyone is welcome to pick it up. Feel free to ping me with questions about the cause or when you'd like a review. -- 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]
