Skip to content

Commit adbb3ce

Browse files
gruuyaclaude
gruuya
and
claude
committed
Prune hash join probe side via a build-side key range bitmap
Map the build side's keys onto a fixed-size bitmap over their range and test container min/max against it, so a scan can skip containers whose values fall in the gaps between keys. The bitmap is bounded at 128 KiB regardless of build side size, and one whose buckets are all set is discarded, so a contiguous key set costs nothing. The key range and distinct count already exist for the bounds predicate, so a dense key set is ruled out before allocating. Integer keys only; other types go unpruned. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
1 parent 714956b commit adbb3ce

10 files changed

Lines changed: 673 additions & 41 deletions

File tree

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

Lines changed: 32 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ use crate::joins::hash_join::stream::{
4040
BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState,
4141
};
4242
use crate::joins::join_hash_map::{JoinHashMapU32, JoinHashMapU64};
43+
use crate::joins::key_range_bitmap::KeyRangeBitmap;
4344
use crate::joins::utils::{
4445
OnceAsync, OnceFut, asymmetric_join_output_partitioning, emits_unmatched_left_rows,
4546
is_existence_join, reorder_output_after_swap, swap_join_projection, update_hash,
@@ -3094,19 +3095,41 @@ async fn collect_left_input(
30943095
.iter()
30953096
.map(|arr| arr.get_array_memory_size())
30963097
.sum::<usize>();
3097-
if left_values.is_empty()
3098-
|| left_values[0].is_empty()
3099-
|| estimated_size > config.optimizer.hash_join_inlist_pushdown_max_size
3100-
|| map.num_of_distinct_key()
3101-
> config
3098+
3099+
let pushdown_inlist = !left_values.is_empty()
3100+
&& !left_values[0].is_empty()
3101+
&& estimated_size <= config.optimizer.hash_join_inlist_pushdown_max_size
3102+
&& map.num_of_distinct_key()
3103+
<= config
31023104
.optimizer
3103-
.hash_join_inlist_pushdown_max_distinct_values
3105+
.hash_join_inlist_pushdown_max_distinct_values;
3106+
3107+
if pushdown_inlist
3108+
&& let Some(in_list_values) = build_struct_inlist_values(&left_values)?
31043109
{
3105-
PushdownStrategy::Map(Arc::clone(&map))
3106-
} else if let Some(in_list_values) = build_struct_inlist_values(&left_values)? {
31073110
PushdownStrategy::InList(in_list_values)
31083111
} else {
3109-
PushdownStrategy::Map(Arc::clone(&map))
3112+
// Past the InList threshold use a bucket bitmap for container pruning.
3113+
let pruning_bitmap = match (left_values.as_slice(), bounds.as_ref()) {
3114+
([keys], Some(bounds)) if !keys.is_empty() => bounds
3115+
.get_column_bounds(0)
3116+
.and_then(|b| {
3117+
KeyRangeBitmap::try_new(
3118+
keys,
3119+
&b.min,
3120+
&b.max,
3121+
map.num_of_distinct_key(),
3122+
)
3123+
})
3124+
.map(Arc::new),
3125+
_ => None,
3126+
};
3127+
if let Some(bitmap) = pruning_bitmap.as_ref() {
3128+
// Held for the join's lifetime, so charge it like the maps.
3129+
reservation.try_grow(bitmap.size())?;
3130+
metrics.build_mem_used.add(bitmap.size());
3131+
}
3132+
PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap)
31103133
}
31113134
};
31123135

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

Lines changed: 60 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ use datafusion_expr_common::dyn_eq::DynHash;
3434
use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef};
3535

3636
use crate::joins::Map;
37+
use crate::joins::key_range_bitmap::KeyRangeBitmap;
3738

3839
/// RandomState wrapper that preserves the seed used to create it.
3940
///
@@ -284,6 +285,8 @@ pub struct HashTableLookupExpr {
284285
map: Arc<Map>,
285286
/// Description for display
286287
description: String,
288+
/// Shared pruning-only view of the build side's key range
289+
pruning_bitmap: Option<Arc<KeyRangeBitmap>>,
287290
}
288291
impl HashTableLookupExpr {
289292
/// Create a new HashTableLookupExpr
@@ -293,6 +296,7 @@ impl HashTableLookupExpr {
293296
/// * `random_state` - SeededRandomState for hashing
294297
/// * `map` - Map to check membership (hash table or array map)
295298
/// * `description` - Description for debugging
299+
/// * `pruning_bitmap` - key-range summary for pruning only, or `None`
296300
///
297301
/// # Public Only for Internal Use:
298302
/// `datafusion-proto` tests require this constructor, but it is not part of
@@ -303,14 +307,21 @@ impl HashTableLookupExpr {
303307
random_state: SeededRandomState,
304308
map: Arc<Map>,
305309
description: String,
310+
pruning_bitmap: Option<Arc<KeyRangeBitmap>>,
306311
) -> Self {
307312
Self {
308313
on_columns,
309314
random_state,
310315
map,
311316
description,
317+
pruning_bitmap,
312318
}
313319
}
320+
321+
/// Which parts of this build side's key range hold keys, for container pruning.
322+
pub fn pruning_bitmap(&self) -> Option<&Arc<KeyRangeBitmap>> {
323+
self.pruning_bitmap.as_ref()
324+
}
314325
}
315326
impl std::fmt::Debug for HashTableLookupExpr {
316327
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -374,12 +385,13 @@ impl PhysicalExpr for HashTableLookupExpr {
374385
self: Arc<Self>,
375386
children: Vec<Arc<dyn PhysicalExpr>>,
376387
) -> Result<Arc<dyn PhysicalExpr>> {
377-
Ok(Arc::new(HashTableLookupExpr::new(
378-
children,
379-
self.random_state.clone(),
380-
Arc::clone(&self.map),
381-
self.description.clone(),
382-
)))
388+
Ok(Arc::new(HashTableLookupExpr {
389+
on_columns: children,
390+
random_state: self.random_state.clone(),
391+
map: Arc::clone(&self.map),
392+
description: self.description.clone(),
393+
pruning_bitmap: self.pruning_bitmap.clone(),
394+
}))
383395
}
384396

385397
fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
@@ -425,6 +437,7 @@ impl PhysicalExpr for HashTableLookupExpr {
425437
random_state: _,
426438
map: _,
427439
description: _,
440+
pruning_bitmap: _,
428441
} = self;
429442

430443
// HashTableLookupExpr holds a runtime Arc<Map> (the build-side hash
@@ -471,6 +484,7 @@ fn evaluate_columns(
471484
mod tests {
472485
use super::*;
473486
use crate::joins::join_hash_map::JoinHashMapU32;
487+
use datafusion_common::ScalarValue;
474488
use datafusion_physical_expr::expressions::Column;
475489
use std::collections::hash_map::DefaultHasher;
476490
use std::hash::Hasher;
@@ -481,6 +495,36 @@ mod tests {
481495
hasher.finish()
482496
}
483497

498+
/// The pruning bitmap must survive `with_new_children`, which a scan calls
499+
/// once per file.
500+
#[test]
501+
fn test_pruning_bitmap_shared_with_derived_children() {
502+
let keys: ArrayRef = Arc::new(arrow::array::Int64Array::from(vec![0, 1_000_000]));
503+
let bitmap = Arc::new(
504+
KeyRangeBitmap::try_new(
505+
&keys,
506+
&ScalarValue::Int64(Some(0)),
507+
&ScalarValue::Int64(Some(1_000_000)),
508+
2,
509+
)
510+
.expect("two keys a million apart leave prunable gaps"),
511+
);
512+
let expr = Arc::new(HashTableLookupExpr::new(
513+
vec![Arc::new(Column::new("a", 0))],
514+
SeededRandomState::with_seed(1),
515+
Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))),
516+
"hash_lookup".to_string(),
517+
Some(Arc::clone(&bitmap)),
518+
));
519+
520+
let derived = expr
521+
.with_new_children(vec![Arc::new(Column::new("a", 7))])
522+
.unwrap();
523+
let derived = derived.downcast_ref::<HashTableLookupExpr>().unwrap();
524+
assert_eq!(derived.children()[0].to_string(), "a@7");
525+
assert!(Arc::ptr_eq(&bitmap, derived.pruning_bitmap().unwrap()));
526+
}
527+
484528
#[test]
485529
fn test_hash_expr_eq_same() {
486530
let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0));
@@ -759,13 +803,15 @@ mod tests {
759803
SeededRandomState::with_seed(1),
760804
Arc::clone(&hash_map),
761805
"lookup".to_string(),
806+
None,
762807
);
763808

764809
let expr2 = HashTableLookupExpr::new(
765810
vec![Arc::clone(&col_a)],
766811
SeededRandomState::with_seed(1),
767812
Arc::clone(&hash_map),
768813
"lookup".to_string(),
814+
None,
769815
);
770816

771817
assert_eq!(expr1, expr2);
@@ -784,13 +830,15 @@ mod tests {
784830
SeededRandomState::with_seed(1),
785831
Arc::clone(&hash_map),
786832
"lookup".to_string(),
833+
None,
787834
);
788835

789836
let expr2 = HashTableLookupExpr::new(
790837
vec![Arc::clone(&col_b)],
791838
SeededRandomState::with_seed(1),
792839
Arc::clone(&hash_map),
793840
"lookup".to_string(),
841+
None,
794842
);
795843

796844
assert_ne!(expr1, expr2);
@@ -807,13 +855,15 @@ mod tests {
807855
SeededRandomState::with_seed(1),
808856
Arc::clone(&hash_map),
809857
"lookup_one".to_string(),
858+
None,
810859
);
811860

812861
let expr2 = HashTableLookupExpr::new(
813862
vec![Arc::clone(&col_a)],
814863
SeededRandomState::with_seed(1),
815864
Arc::clone(&hash_map),
816865
"lookup_two".to_string(),
866+
None,
817867
);
818868

819869
assert_ne!(expr1, expr2);
@@ -833,13 +883,15 @@ mod tests {
833883
SeededRandomState::with_seed(1),
834884
hash_map1,
835885
"lookup".to_string(),
886+
None,
836887
);
837888

838889
let expr2 = HashTableLookupExpr::new(
839890
vec![Arc::clone(&col_a)],
840891
SeededRandomState::with_seed(1),
841892
hash_map2,
842893
"lookup".to_string(),
894+
None,
843895
);
844896

845897
// Different Arc pointers means not equal (uses Arc::ptr_eq)
@@ -857,13 +909,15 @@ mod tests {
857909
SeededRandomState::with_seed(1),
858910
Arc::clone(&hash_map),
859911
"lookup".to_string(),
912+
None,
860913
);
861914

862915
let expr2 = HashTableLookupExpr::new(
863916
vec![Arc::clone(&col_a)],
864917
SeededRandomState::with_seed(1),
865918
Arc::clone(&hash_map),
866919
"lookup".to_string(),
920+
None,
867921
);
868922

869923
// Equal expressions should have equal hashes

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

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ use crate::joins::hash_join::inlist_builder::build_struct_fields;
3131
use crate::joins::hash_join::partitioned_hash_eval::{
3232
HashExpr, HashTableLookupExpr, SeededRandomState,
3333
};
34+
use crate::joins::key_range_bitmap::KeyRangeBitmap;
3435
use crate::repartition::RangeExpr;
3536
use arrow::array::ArrayRef;
3637
use arrow::datatypes::{DataType, Field, Schema};
@@ -136,12 +137,15 @@ fn create_membership_predicate(
136137
)?)))
137138
}
138139
// Use hash table lookup for large build sides
139-
PushdownStrategy::Map(hash_map) => Ok(Some(Arc::new(HashTableLookupExpr::new(
140-
on_right.to_vec(),
141-
random_state.clone(),
142-
hash_map,
143-
"hash_lookup".to_string(),
144-
)) as Arc<dyn PhysicalExpr>)),
140+
PushdownStrategy::Map(hash_map, pruning_bitmap) => {
141+
Ok(Some(Arc::new(HashTableLookupExpr::new(
142+
on_right.to_vec(),
143+
random_state.clone(),
144+
hash_map,
145+
"hash_lookup".to_string(),
146+
pruning_bitmap,
147+
)) as Arc<dyn PhysicalExpr>))
148+
}
145149
// Empty partition - should not create a filter for this
146150
PushdownStrategy::Empty => Ok(None),
147151
}
@@ -277,8 +281,9 @@ pub(crate) struct SharedBuildAccumulator {
277281
pub(crate) enum PushdownStrategy {
278282
/// Use InList for small build sides (< 128MB)
279283
InList(ArrayRef),
280-
/// Use map lookup for large build sides
281-
Map(Arc<Map>),
284+
/// Use map lookup for large build sides. If `Some` the second field
285+
/// represents the build side's keys using a bucket bitmap.
286+
Map(Arc<Map>, Option<Arc<KeyRangeBitmap>>),
282287
/// There was no data in this partition, do not build a dynamic filter for it
283288
Empty,
284289
}

0 commit comments

Comments
 (0)