[fwojciec]

Which issue does this PR close?

• Closes Partitioned hash join keeps its hash maps allocated but unreserved after it finishes (dynamic filter pushdown) #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 each 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 Partitioned hash join keeps its hash maps allocated but unreserved after it finishes (dynamic filter pushdown) #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).
• Added after merging feat(perf): compute routing and lookup hashes once #26249: the hash tables that the routed lookup builds from IN-list partitions (in_list_hash_map) and the merged pruning bitmap are not reserved either. The hash tables grow with the IN lists, which the same limit caps, and the bitmap is at most 128 KiB.

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. Corrected in review: CollectLeft fails after the plan is dropped, with 0 reserved against 136,288 expected, so the test also covers a filter that outlives the plan.
• Added in review: array_map_stays_reserved_with_dynamic_filter, over the same two modes. Its keys are dense, unlike the first test's, so the build map is an array map without a pruning bitmap. After the plan is dropped, pool.reserved() equals the ArrayMap estimate; after the filter is dropped, it is 0. On main both modes fail after the plan is dropped with 0 reserved against 160,000 expected. With only the array_map_key_range argument reverted to the build reservation, this test fails and the other 1,820 joins:: tests pass.
• Added after merging feat(perf): compute routing and lookup hashes once #26249: routed_maps_stay_reserved_with_dynamic_filter runs the Partitioned join over two hash partitions, so the filter is the routed lookup. The maps stay reserved after every stream ends, the amount does not change when the plan is dropped, and it is 0 after the filter is dropped. On main nothing is reserved once the streams end. If the routed lookup holds the maps without their reservation, this test fails and no other joins:: test does.
• The test from Partitioned hash join keeps its hash maps allocated but unreserved after it finishes (dynamic filter pushdown) #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.

[Copilot]

🟢 Approval recommended

No active review issues remain for the changes that retain hash-join memory reservations while dynamic filters hold the maps.

No open issues 🎉

What changed in this PR

Addresses #26980 by keeping hash join build-map memory reserved for as long as dynamic filters retain it. This corrects memory-pool accounting without releasing maps earlier.

Changes:

• Bundles maps with their reservations and shares them through dynamic-filter expressions.
• Preserves the public lookup constructor and includes both reservations in prepared-build accounting.
• Adds lifetime regression coverage for CollectLeft and Partitioned joins.

File	Description
datafusion/​physical-plan/​src/​joins/​mod.rs	Adds ReservedMap to tie map ownership to its memory reservation.
datafusion/​physical-plan/​src/​joins/​hash_join/​shared_bounds.rs	Passes reservation-owning maps into dynamic-filter predicates.
datafusion/​physical-plan/​src/​joins/​hash_join/​partitioned_hash_eval.rs	Retains map reservations while preserving constructor compatibility and map identity comparisons.
datafusion/​physical-plan/​src/​joins/​hash_join/​exec/​prepared.rs	Sums both build reservations and clarifies prepared-build lifetime documentation.
datafusion/​physical-plan/​src/​joins/​hash_join/​exec.rs	Separates map and pruning-bitmap charges from stream-owned memory and tests reservation lifetimes.

_{Review effort: Balanced}

[jayzhan211]

Optional: the PR description says the CollectLeft case of build_map_stays_reserved_with_dynamic_filter passes on main. When I ran it against main, it fails at the assertion after drop(join) (0 reserved vs 136,288 expected). So it also covers a filter that outlives the plan, which might be worth mentioning in the description.

[jayzhan211]

@fwojciec thanks for picking this up! This test only uses sparse keys, so it never reaches the ArrayMap path (array_map_key_range(..., &mut map_reservation)). If I revert just that argument to &mut reservation, all 1,820 joins:: tests still pass, while dense keys lose the charge again (the 61.2 MiB case in #26980). Could you add a dense-key case? This one passes on this branch and fails on main:

/// Dense keys build an `ArrayMap`, which is charged to the map's own
/// reservation in `array_map_key_range` rather than the build's.
#[rstest]
#[tokio::test]
async fn array_map_stays_reserved_with_dynamic_filter(
    #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)]
    mode: PartitionMode,
) -> Result<()> {
    let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)]));
    let batch = |keys: Vec<i64>| {
        let column = Arc::new(Int64Array::from(keys)) as ArrayRef;
        let batch = RecordBatch::try_new(Arc::clone(&schema), vec![column])?;
        TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)
    };
    let on = vec![(
        Arc::new(Column::new_with_schema("k", &schema)?) as _,
        Arc::new(Column::new_with_schema("k", &schema)?) as _,
    )];
    let (join, filter) = hash_join_with_dynamic_filter_and_mode(
        batch((0..20_000).collect())?,
        batch(vec![0, 5, 15])?,
        on,
        JoinType::Inner,
        mode,
    )?;
    let join: Arc<dyn ExecutionPlan> = Arc::new(join);

    let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1 << 22));
    let runtime = RuntimeEnvBuilder::new()
        .with_memory_pool(Arc::clone(&pool))
        .build_arc()?;
    let mut config = SessionConfig::new();
    config
        .options_mut()
        .optimizer
        .hash_join_inlist_pushdown_max_size = 0;
    let task_ctx = Arc::new(
        TaskContext::default()
            .with_runtime(runtime)
            .with_session_config(config),
    );
    let batches = common::collect(join.execute(0, task_ctx)?).await?;
    assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
    assert_phj_used(&join.metrics().unwrap(), true);

    drop(join);
    assert_eq!(
        pool.reserved(),
        ArrayMap::estimate_memory_size(0, 19_999, 20_000)
    );
    drop(filter);
    assert_eq!(pool.reserved(), 0);
    Ok(())
}

[fwojciec]

Thanks, added as-is in 499b6512. Confirmed that it fails in both modes with the array_map_key_range argument reverted.

[fwojciec]

@jayzhan211 Good catch on CollectLeft. My main run used an earlier version of the test. Fixed the description. Appreciate the careful review!

[fwojciec]

Merged main in ccb796d to resolve the conflict with #26249. The routed lookup now holds ReservedMaps as well, so every map the filter keeps stays charged. Also added a test for the routed path in c3db068, since the existing ones only had one build partition.

#26249 also builds hash tables from IN-list partitions and a merged pruning bitmap that aren't reserved. They're bounded, so I left them out of scope next to the IN-list arrays (noted in the description). Happy to reserve them here if you'd prefer.