fix: keep the hash join build map reserved while the dynamic filter holds it

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<Map>,
+    /// Hash table with row indices into `batch`, charged to its own
+    /// reservation because it is also shared with dynamic filters.
+    map: Arc<ReservedMap>,
     /// The input rows for the build side.
     batch: RecordBatch,
     /// Evaluated build-side key expressions, with float `-0.0` rewritten to
@@ -3187,13 +3188,16 @@ fn build_in_memory(
 ) -> Result<JoinLeftData> {
     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,
@@ -3216,7 +3220,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 {
@@ -3313,7 +3317,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
@@ -3357,10 +3361,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());
                 }
@@ -3540,14 +3544,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<usize> {
+    /// 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<i64> {
+        (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<dyn ExecutionPlan>,
+        probe: Arc<dyn ExecutionPlan>,
+        on: JoinOn,
+    }
+
+    fn sparse_key_join_inputs() -> Result<SparseKeyJoinInputs> {
         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;
@@ -3558,8 +3569,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<usize> {
+        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 {
@@ -3605,6 +3630,180 @@ 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<dyn ExecutionPlan> = Arc::new(join);
+
+        let pool: Arc<dyn MemoryPool> = 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::<usize>(), 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::<JoinHashMapU32>())?
+                + keys.len() * size_of::<u32>();
+        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(())
+    }
+
+    /// 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(())
+    }
+
+    /// 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<dyn ExecutionPlan>, 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<dyn ExecutionPlan> = Arc::new(join);
+
+        let pool: Arc<dyn MemoryPool> = 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::<usize>(), 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<PlanProperties>,

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.