From 2426aff7fd6828eb8f4cea79d1079beda90c25d4 Mon Sep 17 00:00:00 2001 From: Filip Wojciechowski Date: Sat, 31 Oct 2026 21:46:36 -0700 Subject: [PATCH 1/3] fix: keep the hash join build map reserved while the dynamic filter holds it --- .../physical-plan/src/joins/hash_join/exec.rs | 133 +++++++++++++++--- .../src/joins/hash_join/exec/prepared.rs | 6 +- .../joins/hash_join/partitioned_hash_eval.rs | 33 ++++- .../src/joins/hash_join/shared_bounds.rs | 6 +- datafusion/physical-plan/src/joins/mod.rs | 54 +++++++ 5 files changed, 201 insertions(+), 31 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 8273a9943845d..36e9bdd6060cb 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -33,7 +33,7 @@ use crate::joins::utils::{ OnceAsync, OnceFut, asymmetric_join_output_partitioning, emits_unmatched_left_rows, is_existence_join, reorder_output_after_swap, swap_join_projection, update_hash, }; -use crate::joins::{JoinOn, JoinOnRef, PartitionMode, SharedBitmapBuilder}; +use crate::joins::{JoinOn, JoinOnRef, PartitionMode, ReservedMap, SharedBitmapBuilder}; use crate::metrics::{Count, MetricBuilder, MetricCategory}; use crate::projection::{ EmbeddedProjection, JoinData, ProjectionExec, try_embed_projection, @@ -191,8 +191,9 @@ fn array_map_key_range( /// Immutable build buffers and their durable reservation. struct JoinBuildData { - /// Hash table with row indices into `batch`, also shared with dynamic filters. - map: Arc, + /// Hash table with row indices into `batch`, charged to its own + /// reservation because it is also shared with dynamic filters. + map: Arc, /// The input rows for the build side. batch: RecordBatch, /// Evaluated build-side key expressions, with float `-0.0` rewritten to @@ -3186,13 +3187,16 @@ fn build_in_memory( ) -> Result { let is_phj_candidate = is_perfect_hash_join_candidate(on_left, schema)?; + // Bound before the map so error paths release the map before its charge. + let mut map_reservation = reservation.new_empty(); + let (join_hash_map, batch, left_values) = if let Some((min_val, max_val)) = array_map_key_range( bounds.as_ref(), schema, batches, on_left, - &mut reservation, + &mut map_reservation, config.execution.perfect_hash_join_small_build_threshold, config.execution.perfect_hash_join_min_key_density, null_equality, @@ -3215,7 +3219,7 @@ fn build_in_memory( // Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX, otherwise use the // `u64` indice variant // Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown - let mut hashmap = new_join_hashmap(num_rows, &mut reservation, metrics)?; + let mut hashmap = new_join_hashmap(num_rows, &mut map_reservation, metrics)?; let scratch_reservation = reservation.new_empty(); if prepared { @@ -3312,7 +3316,7 @@ fn build_in_memory( }) .transpose()?; - let map = Arc::new(join_hash_map); + let map = Arc::new(ReservedMap::new(join_hash_map, map_reservation)); // For an ordinary build, nothing reads the strategy unless this join's own // dynamic filter accumulator exists. A prepared build is shared with joins @@ -3356,10 +3360,10 @@ fn build_in_memory( .map(Arc::new), _ => None, }; - // Held for the join's lifetime, so charge it like the maps; it is - // optional, so skip it rather than fail when the pool is full. + // Shared with the map, so charge it with the map; it is optional, + // so skip it rather than fail when the pool is full. let pruning_bitmap = pruning_bitmap.filter(|bitmap| { - let ok = reservation.try_grow(bitmap.size()).is_ok(); + let ok = map.try_grow(bitmap.size()).is_ok(); if ok { metrics.build_mem_used.add(bitmap.size()); } @@ -3539,14 +3543,21 @@ mod tests { Ok(()) } - /// Runs a join whose 200 keys spread over 2M would size a pruning bitmap - /// at the 128 KiB cap, and reports the bytes the build side reserved. - /// - /// Forces the IN-list threshold off: under `force_hash_collisions` every - /// key hashes to the same bucket, so `num_of_distinct_key()` no longer - /// reflects the real 200 distinct keys and would otherwise take the - /// IN-list branch instead of the one under test. - async fn build_mem_used(limit: usize, dynamic_filters: bool) -> Result { + /// 200 keys spread over 2M: sparse enough for a hash map rather than an + /// array map, and sizing a pruning bitmap at the 128 KiB cap. + fn sparse_keys() -> Vec { + (0..200).map(|i| i * 10_000).collect() + } + + /// A build side of [`sparse_keys`], a probe side with three of them, and + /// the join keys. + struct SparseKeyJoinInputs { + build: Arc, + probe: Arc, + on: JoinOn, + } + + fn sparse_key_join_inputs() -> Result { let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); let batch = |keys: Vec| { let column = Arc::new(Int64Array::from(keys)) as ArrayRef; @@ -3557,8 +3568,22 @@ mod tests { Arc::new(Column::new_with_schema("k", &schema)?) as _, Arc::new(Column::new_with_schema("k", &schema)?) as _, )]; - let build = batch((0..200).map(|i| i * 10_000).collect())?; - let probe = batch(vec![0, 500_000, 1_500_000])?; + Ok(SparseKeyJoinInputs { + build: batch(sparse_keys())?, + probe: batch(vec![0, 500_000, 1_500_000])?, + on, + }) + } + + /// Runs a join over [`sparse_key_join_inputs`] and reports the bytes the + /// build side reserved. + /// + /// Forces the IN-list threshold off: under `force_hash_collisions` every + /// key hashes to the same bucket, so `num_of_distinct_key()` no longer + /// reflects the real 200 distinct keys and would otherwise take the + /// IN-list branch instead of the one under test. + async fn build_mem_used(limit: usize, dynamic_filters: bool) -> Result { + let SparseKeyJoinInputs { build, probe, on } = sparse_key_join_inputs()?; let join = if dynamic_filters { hash_join_with_dynamic_filter(build, probe, on, JoinType::Inner)?.0 } else { @@ -3604,6 +3629,76 @@ mod tests { Ok(()) } + /// The dynamic filter shares the build map and can outlive the join's + /// stream and plan, so the map stays reserved until the last holder drops + /// it, not only until the partition's stream ends. + #[rstest] + #[tokio::test] + async fn build_map_stays_reserved_with_dynamic_filter( + #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)] + mode: PartitionMode, + ) -> Result<()> { + let SparseKeyJoinInputs { build, probe, on } = sparse_key_join_inputs()?; + let (join, filter) = hash_join_with_dynamic_filter_and_mode( + build, + probe, + on, + JoinType::Inner, + mode, + )?; + let join: Arc = Arc::new(join); + + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + 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::(), 3); + + let build_mem_used = join + .metrics() + .unwrap() + .sum_by_name("build_mem_used") + .unwrap() + .as_usize(); + let keys = sparse_keys(); + let map_bytes = + estimate_memory_size::<(u32, u64)>(keys.len(), size_of::())? + + keys.len() * size_of::(); + let keys_array: ArrayRef = Arc::new(Int64Array::from(keys.clone())); + let bitmap_bytes = KeyRangeBitmap::try_new( + &keys_array, + &ScalarValue::Int64(Some(0)), + &ScalarValue::Int64(Some(keys[keys.len() - 1])), + keys.len(), + ) + .map_or(0, |bitmap| bitmap.size()); + let expected = match mode { + // The stream is gone; only the map and its bitmap are still held. + PartitionMode::Partitioned => map_bytes + bitmap_bytes, + PartitionMode::CollectLeft => build_mem_used, + PartitionMode::Auto => unreachable!(), + }; + assert_eq!(pool.reserved(), expected); + + // The filter is the last holder. + drop(join); + assert_eq!(pool.reserved(), map_bytes + bitmap_bytes); + drop(filter); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs index 602d6a5193b06..d20e1431e54df 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs @@ -45,7 +45,7 @@ impl fmt::Debug for PreparedHashJoinBuild { impl PreparedHashJoinBuild { /// Return the retained build reservation, excluding all per-probe state. pub fn reserved_bytes(&self) -> usize { - self.build.reservation.size() + self.build.reservation.size() + self.build.map.reserved_bytes() } /// Return the complete build row count, including duplicate and null keys. @@ -290,8 +290,8 @@ impl HashJoinExecBuilder { /// executions. An existing expression handle is preserved for its probe /// consumers, while the build-report accumulator is reset. /// * Retain a prepared-build lease if a dynamic-filter expression outlives - /// this plan. Its array/map references preserve allocation lifetime, but - /// the lease is what preserves the corresponding memory charge. + /// this plan. The map carries its own charge, but the IN-list arrays + /// share the build batch, and the lease is what preserves their charge. /// * Attach after physical optimizations that rewrite the build child. /// Probe-only rewrites must preserve the attached join's `left()`; replacing /// it or making incompatible key, mode, or join-type changes fails. diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index b580687e158a6..7db81bbd08ae0 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -20,8 +20,8 @@ use datafusion_expr::ColumnarValue; use datafusion_expr_common::dyn_eq::DynHash; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; -use crate::joins::Map; use crate::joins::key_range_bitmap::KeyRangeBitmap; +use crate::joins::{Map, ReservedMap}; /// RandomState wrapper that preserves the seed used to create it. /// @@ -268,8 +268,8 @@ pub struct HashTableLookupExpr { on_columns: Vec, /// Random state for hashing (with seeds preserved for serialization) random_state: SeededRandomState, - /// Map to check against (hash table or array map) - map: Arc, + /// Map to check against (hash table or array map), with its reservation + map: Arc, /// Description for display description: String, /// Shared pruning-only view of the build side's key range @@ -295,6 +295,24 @@ impl HashTableLookupExpr { map: Arc, description: String, pruning_bitmap: Option>, + ) -> Self { + Self::new_reserved( + on_columns, + random_state, + Arc::new(ReservedMap::unreserved(map)), + description, + pruning_bitmap, + ) + } + + /// [`Self::new`] for a join's map, whose memory stays reserved while the + /// lookup holds it. + pub(crate) fn new_reserved( + on_columns: Vec, + random_state: SeededRandomState, + map: Arc, + description: String, + pruning_bitmap: Option>, ) -> Self { Self { on_columns, @@ -329,7 +347,8 @@ impl Hash for HashTableLookupExpr { self.description.hash(state); self.random_state.seed().hash(state); // Hash immutable map identity, consistently with pointer equality below. - Arc::as_ptr(&self.map).hash(state); + let map: &Map = &self.map; + std::ptr::from_ref(map).hash(state); } } @@ -337,10 +356,11 @@ impl PartialEq for HashTableLookupExpr { fn eq(&self, other: &Self) -> bool { // Sharing an immutable map implies equal contents, including across // prepared-build consumers. Distinct maps need no expensive deep comparison. + let (map, other_map): (&Map, &Map) = (&self.map, &other.map); self.on_columns == other.on_columns && self.description == other.description && self.random_state.seed() == other.random_state.seed() - && Arc::ptr_eq(&self.map, &other.map) + && std::ptr::eq(map, other_map) } } @@ -382,7 +402,8 @@ impl PhysicalExpr for HashTableLookupExpr { // Evaluate columns let join_keys = evaluate_columns(&self.on_columns, batch)?; - match self.map.as_ref() { + let map: &Map = &self.map; + match map { Map::HashMap(map) => { with_hashes(&join_keys, self.random_state.random_state(), |hashes| { let array = map.contain_hashes(hashes); diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index c3c37a193e0d8..68a566eb5c8fe 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -11,8 +11,8 @@ use std::sync::Arc; use crate::ExecutionPlan; use crate::ExecutionPlanProperties; use crate::Partitioning; -use crate::joins::Map; use crate::joins::PartitionMode; +use crate::joins::ReservedMap; use crate::joins::hash_join::exec::HASH_JOIN_SEED; use crate::joins::hash_join::inlist_builder::build_struct_fields; use crate::joins::hash_join::not_in::NotIn; @@ -126,7 +126,7 @@ fn create_membership_predicate( } // Use hash table lookup for large build sides PushdownStrategy::Map(hash_map, pruning_bitmap) => { - Ok(Some(Arc::new(HashTableLookupExpr::new( + Ok(Some(Arc::new(HashTableLookupExpr::new_reserved( on_right.to_vec(), random_state.clone(), hash_map, @@ -271,7 +271,7 @@ pub(crate) enum PushdownStrategy { InList(ArrayRef), /// Use map lookup for large build sides. If `Some` the second field /// represents the build side's keys using a bucket bitmap. - Map(Arc, Option>), + Map(Arc, Option>), /// There was no data in this partition, do not build a dynamic filter for it Empty, } diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index a6bc4e63175b7..c4549d847f554 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -6,10 +6,14 @@ use core::fmt; use std::fmt::{Display, Formatter}; +use std::ops::Deref; +use std::sync::Arc; use arrow::array::BooleanBufferBuilder; pub use asof_join::{AsOfJoinExec, AsOfMatchExpr}; pub use cross_join::CrossJoinExec; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::memory_pool::MemoryReservation; use datafusion_physical_expr::PhysicalExprRef; /// # Public Only for Internal Use: /// @@ -94,6 +98,56 @@ impl Map { } } +/// A [`Map`] and the reservation for its memory, which the dynamic filter +/// shares for as long as the plan lives. Held together, the memory stays +/// reserved until the last holder drops the map. +pub(crate) struct ReservedMap { + /// Shared, as [`HashTableLookupExpr::new`] takes one and compares by it. + map: Arc, + /// Also charged for data kept with the map, such as the pruning bitmap. + /// `None` for a map built outside a join; see [`HashTableLookupExpr::new`]. + // Keep the reservation after the map it accounts for. + reservation: Option, +} + +impl ReservedMap { + /// Wraps `map` with the reservation it was charged to. + pub(crate) fn new(map: Map, reservation: MemoryReservation) -> Self { + Self { + map: Arc::new(map), + reservation: Some(reservation), + } + } + + /// Wraps a map that was not charged to a pool. + pub(crate) fn unreserved(map: Arc) -> Self { + Self { + map, + reservation: None, + } + } + + /// Charges `bytes` for data kept with the map. + pub(crate) fn try_grow(&self, bytes: usize) -> Result<()> { + match &self.reservation { + Some(reservation) => reservation.try_grow(bytes), + None => internal_err!("ReservedMap has no reservation to grow"), + } + } + + pub(crate) fn reserved_bytes(&self) -> usize { + self.reservation.as_ref().map_or(0, MemoryReservation::size) + } +} + +impl Deref for ReservedMap { + type Target = Map; + + fn deref(&self) -> &Map { + &self.map + } +} + #[cfg(test)] pub mod test_utils; From 499b6512c643d723b8c31b3735d03a81f071fbbc Mon Sep 17 00:00:00 2001 From: Filip Wojciechowski Date: Tue, 3 Nov 2026 07:28:32 -0800 Subject: [PATCH 2/3] test: add a dense-key case for the ArrayMap reservation Co-authored-by: Jay Zhan --- .../physical-plan/src/joins/hash_join/exec.rs | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 36e9bdd6060cb..aa2d049be8d9c 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3699,6 +3699,61 @@ mod tests { Ok(()) } + /// 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| { + 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 = Arc::new(join); + + let pool: Arc = 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::(), 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(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, From c3db068e05f03ad0144876e0e9de2593e98f4c21 Mon Sep 17 00:00:00 2001 From: Filip Wojciechowski Date: Tue, 3 Nov 2026 08:45:35 -0800 Subject: [PATCH 3/3] test: check that the routed lookup keeps the maps reserved --- .../physical-plan/src/joins/hash_join/exec.rs | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index dcb2b35073d81..5a5d25aee937a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3755,6 +3755,55 @@ mod tests { Ok(()) } + /// With rows in more than one build partition, the filter is a routed lookup + /// that holds every partition's map, so the maps stay reserved until it is + /// dropped. + #[cfg(not(feature = "force_hash_collisions"))] + #[tokio::test] + async fn routed_maps_stay_reserved_with_dynamic_filter() -> Result<()> { + let SparseKeyJoinInputs { build, probe, on } = sparse_key_join_inputs()?; + let repartition = |input: Arc, key: &PhysicalExprRef| { + let partitioning = Partitioning::Hash(vec![Arc::clone(key)], 2); + RepartitionExec::try_new(input, partitioning) + }; + let (join, filter) = hash_join_with_dynamic_filter_and_mode( + Arc::new(repartition(build, &on[0].0)?), + Arc::new(repartition(probe, &on[0].1)?), + on, + JoinType::Inner, + PartitionMode::Partitioned, + )?; + let join: Arc = Arc::new(join); + + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + 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 = crate::execution_plan::collect(Arc::clone(&join), task_ctx).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 3); + assert!(format!("{:?}", filter.current()?).contains("routed by")); + + // Every stream is gone; only the maps and their bitmaps are still held. + let reserved = pool.reserved(); + assert!(reserved > 0); + // The filter is the last holder. + drop(join); + assert_eq!(pool.reserved(), reserved); + drop(filter); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc,