Prune hash join probe side via a build-side key range bitmap

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -40,6 +40,7 @@ use crate::joins::hash_join::stream::{
     BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState,
 };
 use crate::joins::join_hash_map::{JoinHashMapU32, JoinHashMapU64};
+use crate::joins::key_range_bitmap::KeyRangeBitmap;
 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,
@@ -3084,7 +3085,9 @@ async fn collect_left_input(
 
     let map = Arc::new(join_hash_map);
 
-    let membership = if num_rows == 0 {
+    // Nothing reads the strategy unless the dynamic filter accumulator exists,
+    // and that exists only when the pushdown is enabled.
+    let membership = if num_rows == 0 || !should_compute_dynamic_filters {
         PushdownStrategy::Empty
     } else {
         // If the build side is small enough we can use IN list pushdown.
@@ -3094,19 +3097,41 @@ async fn collect_left_input(
             .iter()
             .map(|arr| arr.get_array_memory_size())
             .sum::<usize>();
-        if left_values.is_empty()
-            || left_values[0].is_empty()
-            || estimated_size > config.optimizer.hash_join_inlist_pushdown_max_size
-            || map.num_of_distinct_key()
-                > config
+
+        let pushdown_inlist = !left_values.is_empty()
+            && !left_values[0].is_empty()
+            && estimated_size <= config.optimizer.hash_join_inlist_pushdown_max_size
+            && map.num_of_distinct_key()
+                <= config
                     .optimizer
-                    .hash_join_inlist_pushdown_max_distinct_values
+                    .hash_join_inlist_pushdown_max_distinct_values;
+
+        if pushdown_inlist
+            && let Some(in_list_values) = build_struct_inlist_values(&left_values)?
         {
-            PushdownStrategy::Map(Arc::clone(&map))
-        } else if let Some(in_list_values) = build_struct_inlist_values(&left_values)? {
             PushdownStrategy::InList(in_list_values)
         } else {
-            PushdownStrategy::Map(Arc::clone(&map))
+            // Past the InList threshold use a bucket bitmap for container pruning.
+            let pruning_bitmap = match (left_values.as_slice(), bounds.as_ref()) {
+                ([keys], Some(bounds)) if !keys.is_empty() => bounds
+                    .get_column_bounds(0)
+                    .and_then(|b| {
+                        KeyRangeBitmap::try_new(
+                            keys,
+                            &b.min,
+                            &b.max,
+                            map.num_of_distinct_key(),
+                        )
+                    })
+                    .map(Arc::new),
+                _ => None,
+            };
+            if let Some(bitmap) = pruning_bitmap.as_ref() {
+                // Held for the join's lifetime, so charge it like the maps.
+                reservation.try_grow(bitmap.size())?;
+                metrics.build_mem_used.add(bitmap.size());
+            }
+            PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap)
         }
     };
 
@@ -3264,6 +3289,53 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn no_pruning_state_without_dynamic_filters() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)]));
+        let build = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int64Array::from_iter_values(
+                (0..151).map(|i| i * 10_000),
+            ))],
+        )?;
+        let probe = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int64Array::from_iter_values([
+                0, 500_000, 1_500_000,
+            ]))],
+        )?;
+        let on = vec![(
+            Arc::new(Column::new_with_schema("k", &schema)?) as _,
+            Arc::new(Column::new_with_schema("k", &schema)?) as _,
+        )];
+        let join = join(
+            TestMemoryExec::try_new_exec(&[vec![build]], Arc::clone(&schema), None)?,
+            TestMemoryExec::try_new_exec(&[vec![probe]], Arc::clone(&schema), None)?,
+            on,
+            &JoinType::Inner,
+            NullEquality::NullEqualsNothing,
+        )?;
+
+        // Bounds are still collected to test perfect-hash-join candidacy, so 151
+        // keys over 1.5M would size a bitmap at the 128 KiB cap - far past this.
+        let runtime = RuntimeEnvBuilder::new()
+            .with_memory_limit(100_000, 1.0)
+            .build_arc()?;
+        let mut config = SessionConfig::new();
+        config
+            .options_mut()
+            .optimizer
+            .enable_join_dynamic_filter_pushdown = false;
+        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);
+        Ok(())
+    }
+
     #[derive(Debug)]
     struct PartitionedTestExec {
         cache: Arc<PlanProperties>,

datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs

@@ -34,6 +34,7 @@ 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;
 
 /// RandomState wrapper that preserves the seed used to create it.
 ///
@@ -284,6 +285,8 @@ pub struct HashTableLookupExpr {
     map: Arc<Map>,
     /// Description for display
     description: String,
+    /// Shared pruning-only view of the build side's key range
+    pruning_bitmap: Option<Arc<KeyRangeBitmap>>,
 }
 impl HashTableLookupExpr {
     /// Create a new HashTableLookupExpr
@@ -293,6 +296,7 @@ impl HashTableLookupExpr {
     /// * `random_state` - SeededRandomState for hashing
     /// * `map` - Map to check membership (hash table or array map)
     /// * `description` - Description for debugging
+    /// * `pruning_bitmap` - key-range summary for pruning only, or `None`
     ///
     /// # Public Only for Internal Use:
     /// `datafusion-proto` tests require this constructor, but it is not part of
@@ -303,14 +307,21 @@ impl HashTableLookupExpr {
         random_state: SeededRandomState,
         map: Arc<Map>,
         description: String,
+        pruning_bitmap: Option<Arc<KeyRangeBitmap>>,
     ) -> Self {
         Self {
             on_columns,
             random_state,
             map,
             description,
+            pruning_bitmap,
         }
     }
+
+    /// Which parts of this build side's key range hold keys, for container pruning.
+    pub fn pruning_bitmap(&self) -> Option<&Arc<KeyRangeBitmap>> {
+        self.pruning_bitmap.as_ref()
+    }
 }
 impl std::fmt::Debug for HashTableLookupExpr {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -374,12 +385,13 @@ impl PhysicalExpr for HashTableLookupExpr {
         self: Arc<Self>,
         children: Vec<Arc<dyn PhysicalExpr>>,
     ) -> Result<Arc<dyn PhysicalExpr>> {
-        Ok(Arc::new(HashTableLookupExpr::new(
-            children,
-            self.random_state.clone(),
-            Arc::clone(&self.map),
-            self.description.clone(),
-        )))
+        Ok(Arc::new(HashTableLookupExpr {
+            on_columns: children,
+            random_state: self.random_state.clone(),
+            map: Arc::clone(&self.map),
+            description: self.description.clone(),
+            pruning_bitmap: self.pruning_bitmap.clone(),
+        }))
     }
 
     fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
@@ -425,6 +437,7 @@ impl PhysicalExpr for HashTableLookupExpr {
             random_state: _,
             map: _,
             description: _,
+            pruning_bitmap: _,
         } = self;
 
         // HashTableLookupExpr holds a runtime Arc<Map> (the build-side hash
@@ -471,6 +484,7 @@ fn evaluate_columns(
 mod tests {
     use super::*;
     use crate::joins::join_hash_map::JoinHashMapU32;
+    use datafusion_common::ScalarValue;
     use datafusion_physical_expr::expressions::Column;
     use std::collections::hash_map::DefaultHasher;
     use std::hash::Hasher;
@@ -481,6 +495,36 @@ mod tests {
         hasher.finish()
     }
 
+    /// The pruning bitmap must survive `with_new_children`, which a scan calls
+    /// once per file.
+    #[test]
+    fn test_pruning_bitmap_shared_with_derived_children() {
+        let keys: ArrayRef = Arc::new(arrow::array::Int64Array::from(vec![0, 1_000_000]));
+        let bitmap = Arc::new(
+            KeyRangeBitmap::try_new(
+                &keys,
+                &ScalarValue::Int64(Some(0)),
+                &ScalarValue::Int64(Some(1_000_000)),
+                2,
+            )
+            .expect("two keys a million apart leave prunable gaps"),
+        );
+        let expr = Arc::new(HashTableLookupExpr::new(
+            vec![Arc::new(Column::new("a", 0))],
+            SeededRandomState::with_seed(1),
+            Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))),
+            "hash_lookup".to_string(),
+            Some(Arc::clone(&bitmap)),
+        ));
+
+        let derived = expr
+            .with_new_children(vec![Arc::new(Column::new("a", 7))])
+            .unwrap();
+        let derived = derived.downcast_ref::<HashTableLookupExpr>().unwrap();
+        assert_eq!(derived.children()[0].to_string(), "a@7");
+        assert!(Arc::ptr_eq(&bitmap, derived.pruning_bitmap().unwrap()));
+    }
+
     #[test]
     fn test_hash_expr_eq_same() {
         let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0));
@@ -759,13 +803,15 @@ mod tests {
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup".to_string(),
+            None,
         );
 
         let expr2 = HashTableLookupExpr::new(
             vec![Arc::clone(&col_a)],
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup".to_string(),
+            None,
         );
 
         assert_eq!(expr1, expr2);
@@ -784,13 +830,15 @@ mod tests {
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup".to_string(),
+            None,
         );
 
         let expr2 = HashTableLookupExpr::new(
             vec![Arc::clone(&col_b)],
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup".to_string(),
+            None,
         );
 
         assert_ne!(expr1, expr2);
@@ -807,13 +855,15 @@ mod tests {
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup_one".to_string(),
+            None,
         );
 
         let expr2 = HashTableLookupExpr::new(
             vec![Arc::clone(&col_a)],
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup_two".to_string(),
+            None,
         );
 
         assert_ne!(expr1, expr2);
@@ -833,13 +883,15 @@ mod tests {
             SeededRandomState::with_seed(1),
             hash_map1,
             "lookup".to_string(),
+            None,
         );
 
         let expr2 = HashTableLookupExpr::new(
             vec![Arc::clone(&col_a)],
             SeededRandomState::with_seed(1),
             hash_map2,
             "lookup".to_string(),
+            None,
         );
 
         // Different Arc pointers means not equal (uses Arc::ptr_eq)
@@ -857,13 +909,15 @@ mod tests {
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup".to_string(),
+            None,
         );
 
         let expr2 = HashTableLookupExpr::new(
             vec![Arc::clone(&col_a)],
             SeededRandomState::with_seed(1),
             Arc::clone(&hash_map),
             "lookup".to_string(),
+            None,
         );
 
         // Equal expressions should have equal hashes

datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs

@@ -31,6 +31,7 @@ use crate::joins::hash_join::inlist_builder::build_struct_fields;
 use crate::joins::hash_join::partitioned_hash_eval::{
     HashExpr, HashTableLookupExpr, SeededRandomState,
 };
+use crate::joins::key_range_bitmap::KeyRangeBitmap;
 use crate::repartition::RangeExpr;
 use arrow::array::ArrayRef;
 use arrow::datatypes::{DataType, Field, Schema};
@@ -136,12 +137,15 @@ fn create_membership_predicate(
             )?)))
         }
         // Use hash table lookup for large build sides
-        PushdownStrategy::Map(hash_map) => Ok(Some(Arc::new(HashTableLookupExpr::new(
-            on_right.to_vec(),
-            random_state.clone(),
-            hash_map,
-            "hash_lookup".to_string(),
-        )) as Arc<dyn PhysicalExpr>)),
+        PushdownStrategy::Map(hash_map, pruning_bitmap) => {
+            Ok(Some(Arc::new(HashTableLookupExpr::new(
+                on_right.to_vec(),
+                random_state.clone(),
+                hash_map,
+                "hash_lookup".to_string(),
+                pruning_bitmap,
+            )) as Arc<dyn PhysicalExpr>))
+        }
         // Empty partition - should not create a filter for this
         PushdownStrategy::Empty => Ok(None),
     }
@@ -277,8 +281,9 @@ pub(crate) struct SharedBuildAccumulator {
 pub(crate) enum PushdownStrategy {
     /// Use InList for small build sides (< 128MB)
     InList(ArrayRef),
-    /// Use map lookup for large build sides
-    Map(Arc<Map>),
+    /// Use map lookup for large build sides. If `Some` the second field
+    /// represents the build side's keys using a bucket bitmap.
+    Map(Arc<Map>, Option<Arc<KeyRangeBitmap>>),
     /// There was no data in this partition, do not build a dynamic filter for it
     Empty,
 }