From 0e6e9c118bd46f637790cd0575ac03b397177f67 Mon Sep 17 00:00:00 2001 From: Marko Grujic Date: Mon, 21 Sep 2026 14:38:19 +0200 Subject: [PATCH 1/6] 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) --- .../physical-plan/src/joins/hash_join/exec.rs | 41 ++- .../joins/hash_join/partitioned_hash_eval.rs | 66 ++++- .../src/joins/hash_join/shared_bounds.rs | 21 +- .../src/joins/key_range_bitmap.rs | 238 ++++++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 1 + datafusion/proto/tests/cases/plans/exprs.rs | 1 + .../pruning/src/key_range_bitmap_expr.rs | 103 ++++++++ datafusion/pruning/src/lib.rs | 1 + datafusion/pruning/src/pruning_predicate.rs | 215 ++++++++++++++-- .../test_files/dynamic_row_group_pruning.slt | 27 ++ 10 files changed, 673 insertions(+), 41 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/key_range_bitmap.rs create mode 100644 datafusion/pruning/src/key_range_bitmap_expr.rs diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 1bffdb1ab2d7c..e2e743985d678 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/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, @@ -3323,19 +3324,41 @@ async fn collect_left_input( .iter() .map(|arr| arr.get_array_memory_size()) .sum::(); - 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) } }; 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 0fa5615db3389..921a38970562e 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 @@ -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, /// Description for display description: String, + /// Shared pruning-only view of the build side's key range + pruning_bitmap: Option>, } 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, description: String, + pruning_bitmap: Option>, ) -> 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> { + self.pruning_bitmap.as_ref() + } } impl std::fmt::Debug for HashTableLookupExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -363,12 +374,13 @@ impl PhysicalExpr for HashTableLookupExpr { self: Arc, children: Vec>, ) -> Result> { - 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 { @@ -414,6 +426,7 @@ impl PhysicalExpr for HashTableLookupExpr { random_state: _, map: _, description: _, + pruning_bitmap: _, } = self; // HashTableLookupExpr holds a runtime Arc (the build-side hash @@ -460,6 +473,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; @@ -470,6 +484,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::().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)); @@ -748,6 +792,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -755,6 +800,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_eq!(expr1, expr2); @@ -773,6 +819,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -780,6 +827,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -796,6 +844,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_one".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -803,6 +852,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_two".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -822,6 +872,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map1, "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -829,6 +880,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map2, "lookup".to_string(), + None, ); // Different Arc pointers means not equal (uses Arc::ptr_eq) @@ -846,6 +898,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -853,6 +906,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); // Equal expressions should have equal 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 62087c14c5179..b6187f10906b1 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/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)), + 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)) + } // 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), + /// 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>), /// There was no data in this partition, do not build a dynamic filter for it Empty, } diff --git a/datafusion/physical-plan/src/joins/key_range_bitmap.rs b/datafusion/physical-plan/src/joins/key_range_bitmap.rs new file mode 100644 index 0000000000000..efdb5467c897e --- /dev/null +++ b/datafusion/physical-plan/src/joins/key_range_bitmap.rs @@ -0,0 +1,238 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Bitmap over a hash join build side's key range, for container pruning. + +use arrow::array::{Array, ArrayRef, BooleanArray, downcast_integer_array}; +use arrow::buffer::{BooleanBuffer, MutableBuffer}; +use arrow::util::bit_util; +use datafusion_common::ScalarValue; +use num_traits::AsPrimitive; + +use crate::joins::array_map::ArrayMap; + +/// Largest bitmap, in bits (128 KiB), which bounds what one build side can hold. +const MAX_BUCKETS: u64 = 1 << 20; + +/// Which parts of a build side's key range hold at least one key. +/// +/// Keys map order-preservingly onto buckets spanning `[min, max]`, one bit each, +/// so "can a container holding `[lo, hi]` match?" is a scan of the bits spanning +/// `[lo, hi]`: all clear proves no build key lies in it. +/// +/// Below, keys 3 to 97 over eight buckets: `offset` is 3, `span` is 97 - 3, and +/// eight buckets need `shift` = 4, so each covers `1 << 4` keys starting at +/// `offset`. Only buckets 0 to `span >> shift` can hold a key. Real bitmaps run +/// to `MAX_BUCKETS`; eight is what fits on a line. +/// +/// ```text +/// keys 3 12 42 97 +/// v v v v +/// bucket | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | +/// covers 3-18 19-34 35-50 51-66 67-82 83-98 unused unused +/// bit 1 0 1 0 0 1 0 0 +/// ``` +/// +/// A container is kept when every bucket its `[min, max]` touches has a bit set. +#[derive(Debug)] +pub struct KeyRangeBitmap { + /// Smallest key, in the `u64` ordering the join's array map uses. + offset: u64, + /// `max - min`: a key is in range iff `key.wrapping_sub(offset) <= span`. + span: u64, + /// Each bucket covers `1 << shift` keys. + shift: u32, + /// One bit per bucket. + bits: BooleanBuffer, +} + +impl KeyRangeBitmap { + /// Maps `array`'s non-null values, which must lie within `[min, max]`. + /// + /// `None` when the key type has no `u64` ordering, or when the bitmap could + /// not exclude anything the `[min, max]` bounds predicate does not already. + pub fn try_new( + array: &ArrayRef, + min: &ScalarValue, + max: &ScalarValue, + distinct_keys: usize, + ) -> Option { + if !ArrayMap::is_supported_type(array.data_type()) { + return None; + } + let offset = ArrayMap::key_to_u64(min)?; + let span = ArrayMap::key_to_u64(max)?.wrapping_sub(offset); + + // One bucket per key value where the range allows it, so the bitmap is + // exact; otherwise as many as the cap permits. + let buckets = span.saturating_add(1).min(MAX_BUCKETS).next_power_of_two(); + + // Use a heuristic to bail-out for the common shape which renders each bucket + // set (contiguous surrogate keys, a date range) before actually computing it. + if (distinct_keys as u64).saturating_mul(2) > span.min(buckets) { + return None; + } + + let shift = (0u32..63).find(|s| (span >> s) < buckets).unwrap_or(63); + + let bits = bucket_bits(array, offset, shift, buckets); + + // Every bucket set means nothing can be excluded. + let reachable = ((span >> shift) + 1) as usize; + bits.slice(0, reachable).has_false().then_some(Self { + offset, + span, + shift, + bits, + }) + } + + /// One Boolean per `[min[i], max[i]]` container interval: `false` proves no + /// build key lies in that container, `None` where a bound is absent or the + /// statistics are not integer keys. + pub fn may_contain_ranges(&self, min: &dyn Array, max: &dyn Array) -> BooleanArray { + // Both bounds must be the same integer type for the tuple pattern to match. + downcast_integer_array!( + (min, max) => { + min.iter() + .zip(max.iter()) + .map(|bounds| match bounds { + (Some(lo), Some(hi)) => { + Some(self.may_contain_range(lo.as_(), hi.as_())) + } + _ => None, + }) + .collect() + } + _ => BooleanArray::new_null(min.len()), + ) + } + + /// Might a container spanning `[lo, hi]` hold a build key? `false` proves it + /// cannot. Bounds use the same `u64` key ordering as the join's array map. + pub fn may_contain_range(&self, lo: u64, hi: u64) -> bool { + // Offsets from the minimum; anything outside `[min, max]` wraps to + // something huge, which the clamping below turns into "past the end". + let lo_off = lo.wrapping_sub(self.offset); + let hi_off = hi.wrapping_sub(self.offset); + if lo_off > self.span && hi_off > self.span && lo_off <= hi_off { + return false; + } + let lo_b = if lo_off > self.span { + 0 + } else { + lo_off >> self.shift + }; + let hi_b = (hi_off.min(self.span) >> self.shift).max(lo_b); + self.bits + .slice(lo_b as usize, (hi_b - lo_b + 1) as usize) + .has_true() + } + + /// Bytes held, for the build-side memory reservation. + pub fn size(&self) -> usize { + self.bits.inner().capacity() + } + + /// Buckets set and buckets total, for `EXPLAIN`. + pub fn fill_stats(&self) -> (usize, usize) { + (self.bits.count_set_bits(), self.bits.len()) + } +} + +/// One bit per bucket, set for each bucket holding a non-null key of `keys`. +fn bucket_bits(keys: &dyn Array, offset: u64, shift: u32, buckets: u64) -> BooleanBuffer { + let mut words = MutableBuffer::new_null(buckets as usize); + downcast_integer_array!( + keys => { + let bytes = words.as_slice_mut(); + let mut set = |key: u64| { + let bucket = (key.wrapping_sub(offset) >> shift).min(buckets - 1); + bit_util::set_bit(bytes, bucket as usize); + }; + if keys.null_count() == 0 { + keys.values().iter().for_each(|key| set((*key).as_())); + } else { + keys.iter().flatten().for_each(|key| set(key.as_())); + } + } + _ => unreachable!("guarded by ArrayMap::is_supported_type"), + ); + BooleanBuffer::new(words.into(), 0, buckets as usize) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use std::sync::Arc; + + fn bitmap(keys: &[i64]) -> Option { + let a: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); + KeyRangeBitmap::try_new( + &a, + &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), + &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), + keys.iter().collect::>().len(), + ) + } + fn may(b: &KeyRangeBitmap, lo: i64, hi: i64) -> bool { + b.may_contain_range(lo as u64, hi as u64) + } + + #[test] + fn prunes_gaps_and_never_drops_a_match() { + let keys: Vec = (0..2_000_000).filter(|k| k % 10000 < 200).collect(); + let s = bitmap(&keys).expect("not saturated"); + let mut kept = 0; + for i in 0..2000i64 { + let (lo, hi) = (1000 * i, 1000 * i + 999); + let truth = keys + .binary_search(&lo) + .map_or_else(|p| keys.get(p).is_some_and(|k| *k <= hi), |_| true); + let got = may(&s, lo, hi); + assert!( + !truth || got, + "dropped a container holding a key: {lo}..{hi}" + ); + kept += got as usize; + } + assert!(kept <= 220, "expected ~200 kept, got {kept}"); + } + + #[test] + fn contiguous_keys_no_bitmap() { + assert!(bitmap(&(0..5000i64).collect::>()).is_none()); + } + + #[test] + fn test_bitmap_range_probing() { + let s = bitmap(&[0, 1_000_000]).expect("not saturated"); + assert!(!may(&s, -500, -10)); + assert!(!may(&s, 2_000_000, 3_000_000)); + assert!(may(&s, -500, 10)); + + let s = bitmap(&[-1_000_000, -5, 0, 5, 1_000_000]).expect("not saturated"); + assert!(may(&s, -10, 10)); + assert!(!may(&s, -900_000, -800_000)); + + // The widest possible span must not overflow the bucket sizing. + let s = bitmap(&[i64::MIN, 0, i64::MAX]).expect("not saturated"); + assert!(may(&s, -1, 1)); + assert!(!may(&s, 1 << 50, 1 << 51)); + } +} diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 0d2480c131469..43ab24d87d128 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -55,6 +55,7 @@ mod join_filter; /// Note: This module is public for internal testing purposes only /// and is not guaranteed to be stable across versions. pub mod join_hash_map; +pub mod key_range_bitmap; use array_map::ArrayMap; use utils::JoinHashMapType; diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs index ee6f9819a3166..a7f6aa72000d9 100644 --- a/datafusion/proto/tests/cases/plans/exprs.rs +++ b/datafusion/proto/tests/cases/plans/exprs.rs @@ -117,6 +117,7 @@ fn roundtrip_hash_table_lookup_expr_to_lit() -> Result<()> { datafusion::physical_plan::joins::SeededRandomState::with_seed(0), hash_map, "test_lookup".to_string(), + None, )); // Create a filter with the lookup expression diff --git a/datafusion/pruning/src/key_range_bitmap_expr.rs b/datafusion/pruning/src/key_range_bitmap_expr.rs new file mode 100644 index 0000000000000..ba30a052cac10 --- /dev/null +++ b/datafusion/pruning/src/key_range_bitmap_expr.rs @@ -0,0 +1,103 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; +use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; + +/// Tests containers' `[min, max]` statistics against a build side's +/// [`KeyRangeBitmap`]: one nullable Boolean per container, where `false` proves +/// no build key lies in it and `NULL` means the bounds could not decide. +#[derive(Debug)] +pub(crate) struct KeyRangeBitmapPruningExpr { + pub(crate) min: PhysicalExprRef, + pub(crate) max: PhysicalExprRef, + pub(crate) bitmap: Arc, +} + +impl PartialEq for KeyRangeBitmapPruningExpr { + fn eq(&self, other: &Self) -> bool { + self.min.eq(&other.min) + && self.max.eq(&other.max) + && Arc::ptr_eq(&self.bitmap, &other.bitmap) + } +} +impl Eq for KeyRangeBitmapPruningExpr {} + +impl Hash for KeyRangeBitmapPruningExpr { + fn hash(&self, state: &mut H) { + self.min.hash(state); + self.max.hash(state); + Arc::as_ptr(&self.bitmap).hash(state); + } +} + +impl Display for KeyRangeBitmapPruningExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let (set, total) = self.bitmap.fill_stats(); + write!( + f, + "KEY_RANGE_BITMAP({}, {}, {set}/{total} buckets)", + self.min, self.max + ) + } +} + +impl PhysicalExpr for KeyRangeBitmapPruningExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let rows = batch.num_rows(); + let min = self.min.evaluate(batch)?.into_array(rows)?; + let max = self.max.evaluate(batch)?.into_array(rows)?; + let matches = self.bitmap.may_contain_ranges(&min, &max); + Ok(ColumnarValue::Array(Arc::new(matches))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + bitmap: Arc::clone(&self.bitmap), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +} diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index 88a7fdda5e733..3e7c2ad7aa7ab 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -19,6 +19,7 @@ mod file_pruner; mod in_list; +mod key_range_bitmap_expr; mod primitive_in_list; mod pruning_predicate; mod string_in_list; diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index b49b72058e0cd..49725983a391b 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -23,6 +23,7 @@ use std::collections::HashSet; use std::sync::Arc; use crate::in_list::{SetMembership, unwrap_scalar}; +use crate::key_range_bitmap_expr::KeyRangeBitmapPruningExpr; use crate::primitive_in_list::PrimitiveInListDomain; use crate::string_in_list::{BinaryInListPruningExpr, StringInListPruningExpr}; @@ -52,6 +53,7 @@ use datafusion_expr_common::operator::Operator; use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee}; use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr}; use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt; +use datafusion_physical_plan::joins::HashTableLookupExpr; use datafusion_physical_plan::{ColumnarValue, PhysicalExpr}; /// Used to prove that arbitrary predicates (boolean expression) can not @@ -1520,6 +1522,62 @@ impl CompactInListDomain { } } +/// Statistics expressions for `column`: its min, its max, and "not every row is +/// NULL". Appends to `required_columns`, rolling the appends back if any part +/// cannot be built, since [`RequiredColumns::stat_column_expr`] only appends. +fn min_max_non_null_exprs( + column: &phys_expr::Column, + column_expr: &Arc, + field: &Field, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option<(PhysicalExprRef, PhysicalExprRef, PhysicalExprRef)> { + let appended = required_columns.columns.len(); + let statistics = (|| { + let min = required_columns + .min_column_expr(column, column_expr, field) + .ok()?; + let max = required_columns + .max_column_expr(column, column_expr, field) + .ok()?; + let non_null = + build_is_null_column_expr(column_expr, schema, required_columns, true)?; + Some((min, max, non_null)) + })(); + if statistics.is_none() { + required_columns.columns.truncate(appended); + } + statistics +} + +/// Builds a "may match" expression for a [`HashTableLookupExpr`] (a large join +/// build side pushed down as an opaque hash-table lookup): tests container +/// min/max stats against the build side's key-range bitmap, which is fixed size +/// and built once with the join's hash table. +fn build_hash_lookup_pruning_expr( + lookup: &HashTableLookupExpr, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option> { + let bitmap = Arc::clone(lookup.pruning_bitmap()?); + let on_columns = lookup.children(); + let [column_expr] = on_columns[..] else { + return None; + }; + let column = column_expr.downcast_ref::()?; + let field = schema.fields().get(column.index())?; + if field.name() != column.name() { + return None; + } + let (min, max, non_null) = + min_max_non_null_exprs(column, column_expr, field, schema, required_columns)?; + Some(Arc::new(phys_expr::BinaryExpr::new( + non_null, + Operator::And, + Arc::new(KeyRangeBitmapPruningExpr { min, max, bitmap }), + ))) +} + /// Keep large literal lists of supported ordered types compact instead of /// building a per-value tree: an OR tree for `IN`, an AND chain for `NOT IN`. /// @@ -1596,24 +1654,8 @@ fn build_compact_in_list_expr( Some(false), )))); } - // Roll back appended statistics columns if the compact rewrite cannot be - // completed. `RequiredColumns::stat_column_expr` only appends entries. - let required_columns_len = required_columns.columns.len(); - let statistics = (|| { - let min = required_columns - .min_column_expr(column, in_list.expr(), field) - .ok()?; - let max = required_columns - .max_column_expr(column, in_list.expr(), field) - .ok()?; - let non_null = - build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; - Some((min, max, non_null)) - })(); - let Some((min, max, non_null)) = statistics else { - required_columns.columns.truncate(required_columns_len); - return None; - }; + let (min, max, non_null) = + min_max_non_null_exprs(column, in_list.expr(), field, schema, required_columns)?; let may_match = match domain { CompactInListDomain::String(values) => { Arc::new(StringInListPruningExpr::new(membership, min, max, values)) @@ -1816,6 +1858,41 @@ fn build_predicate_expression( return unhandled_hook.handle(expr); } } + if let Some(lookup) = expr.downcast_ref::() { + return build_hash_lookup_pruning_expr(lookup, schema, required_columns) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } + // A partitioned hash join hides its per-partition filters under a `CASE` on the + // repartition hash. A row takes exactly one branch, so a container may match + // only if some branch may: the branches' disjunction is a sound relaxation, and + // the `WHEN`s (a hash, which no statistics describe) can be dropped. + if let Some(case) = expr.downcast_ref::() { + // Only a Boolean `CASE` is a predicate; anything else is a value for + // whatever compares it to handle. + if !matches!(case.data_type(schema), Ok(DataType::Boolean)) { + return unhandled_hook.handle(expr); + } + // A missing `ELSE` yields NULL, which never matches, so it adds nothing. + return case + .when_then_expr() + .iter() + .map(|(_, then)| then) + .chain(case.else_expr()) + .map(|branch| { + build_predicate_expression( + branch, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + properties, + ) + }) + .reduce(|acc, branch| { + Arc::new(phys_expr::BinaryExpr::new(acc, Operator::Or, branch)) as _ + }) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } let (left, op, right) = { if let Some(bin_expr) = expr.downcast_ref::() { @@ -2407,6 +2484,9 @@ mod tests { self as phys_expr, DynamicFilterPhysicalExpr, }; use datafusion_physical_expr::planner::logical2physical; + use datafusion_physical_plan::joins::join_hash_map::JoinHashMapU32; + use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; + use datafusion_physical_plan::joins::{Map, SeededRandomState}; use itertools::Itertools; #[derive(Debug, Default)] @@ -7340,4 +7420,103 @@ mod tests { "c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1"; assert_eq!(res.to_string(), expected); } + + /// A hash-table lookup over `column` carrying a bitmap over `keys`. + fn hash_lookup( + column: &Arc, + keys: &[i64], + ) -> Arc { + let array: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); + let bitmap = KeyRangeBitmap::try_new( + &array, + &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), + &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), + keys.len(), + ); + Arc::new(HashTableLookupExpr::new( + vec![Arc::clone(column)], + SeededRandomState::with_seed(1), + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(1)))), + "hash_lookup".to_string(), + bitmap.map(Arc::new), + )) + } + + #[test] + fn test_hash_lookup_pruning_via_min_max() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); + let column: Arc = Arc::new(phys_expr::Column::new("b", 0)); + // Gapped keys: 10,20,30 then a far-away 100000, so the bitmap is sparse. + let lookup = hash_lookup(&column, &[10, 20, 30, 100_000]); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(lookup) + .unwrap(); + + // An opaque lookup yields no literal guarantees, so `contained()` is never + // consulted and every exclusion below comes from the min/max rewrite. + assert!(predicate.literal_guarantees().is_empty()); + + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i64( + vec![Some(5000), Some(15), Some(50000)], + vec![Some(8000), Some(25), Some(60000)], + ), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 + // holds 20. The bitmap excludes the first and third. + assert_eq!(result, vec![false, true, false]); + } + + #[test] + fn test_partition_routed_hash_lookup_pruning() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); + let column: Arc = Arc::new(phys_expr::Column::new("b", 0)); + let literal = |value: ScalarValue| -> Arc { + Arc::new(phys_expr::Literal::new(value)) + }; + // One branch per build partition, each holding its own slice of the build + // side, as `build_partitioned_filter` produces them. + let case: Arc = Arc::new( + phys_expr::CaseExpr::try_new( + Some(literal(ScalarValue::UInt64(Some(0)))), + vec![ + ( + literal(ScalarValue::UInt64(Some(0))), + hash_lookup(&column, &[10, 20, 30]), + ), + ( + literal(ScalarValue::UInt64(Some(1))), + hash_lookup(&column, &[40, 50, 60]), + ), + ], + // Partitions with no build rows reject everything routed to them. + Some(literal(ScalarValue::Boolean(Some(false)))), + ) + .unwrap(), + ); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(case) + .unwrap(); + + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i32( + vec![Some(5), Some(15), Some(45), Some(100)], + vec![Some(8), Some(25), Some(55), Some(200)], + ), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Container 1 ([15,25]) holds 20 from the first branch and container 2 + // ([45,55]) holds 50 from the second, so a branch may match in each. The + // other two intersect neither branch, nor the `ELSE`. + assert_eq!(result, vec![false, true, true, false]); + } } diff --git a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt index c674dede75706..fec01855094f7 100644 --- a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt +++ b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt @@ -318,6 +318,33 @@ SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5; statement ok set datafusion.execution.parquet.pushdown_filters = true; +# Regression test: a hash join's build side must get the same compact, +# sorted-domain min/max pruning an ordinary large `IN (...)` list already +# gets - and do better than the plain min/max *bounds* check every join +# already gets for free. `dim`'s overall envelope [1000, 3050] spans RG 1 +# (b=2000..2099) entirely, so bounds alone cannot exclude it; only the +# discrete check can prove none of {1000, 1050, 3050} falls inside it. +# Force the hash-table-lookup path (rather than `InList`) regardless of size. +statement ok +CREATE TABLE dim AS VALUES (1000), (1050), (3050); + +statement ok +set datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values = 0; + +query TT +explain analyze select rgsel.b from dim join rgsel on dim.column1 = rgsel.b; +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(column1@0, b@0)], projection=[b@1], metrics=[output_rows=3, elapsed_compute=, output_bytes=, output_batches=1, build_mem_used=, array_map_created_count=0, build_input_batches=1, build_input_rows=3, input_batches=2, input_rows=3, build_time=, join_time=, avg_fanout=100% (3/3), probe_hit_rate=100% (3/3)] +02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_row_group_pruning/rgsel.parquet]]}, projection=[b], file_type=parquet, predicate=DynamicFilter [ b@1 >= 1000 AND b@1 <= 3050 AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=b_null_count@1 != row_count@2 AND b_max@0 >= 1000 AND b_null_count@1 != row_count@2 AND b_min@3 <= 3050 AND b_null_count@1 != row_count@2 AND KEY_RANGE_BITMAP(b_min@3, b_max@0, 3/4096 buckets), required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, output_batches=2, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 2 matched, row_groups_pruned_bloom_filter=2 total → 2 matched, page_index_pages_pruned=20 total → 3 matched, page_index_rows_pruned=200 total → 30 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, bytes_processed=, bytes_scanned=, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=27, row_groups_pruned_dynamic_filter=0, predicate_cache_inner_records=200, predicate_cache_records=13, bloom_filter_eval_time=, metadata_load_time=, page_index_eval_time=, row_pushdown_eval_time=, statistics_eval_time=, time_elapsed_opening=, time_elapsed_processing=, time_elapsed_scanning_total=, time_elapsed_scanning_until_data=, output_rows_skew=, scan_efficiency_ratio=] + +statement ok +RESET datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values; + +statement ok +drop table dim; + statement ok drop table rgsel; From 9a684f036ce9705a37d287e95afcbc067596959c Mon Sep 17 00:00:00 2001 From: Marko Grujic Date: Wed, 23 Sep 2026 08:18:16 +0200 Subject: [PATCH 2/6] Fix CASE full-match inference and skip unused pruning state Co-Authored-By: Claude Opus 5 (1M context) --- .../physical-plan/src/joins/hash_join/exec.rs | 51 ++++++++++++++++++- datafusion/pruning/src/pruning_predicate.rs | 6 ++- .../test_files/parquet_filter_pushdown.slt | 47 +++++++++++++++++ 3 files changed, 102 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index e2e743985d678..3a9267160507f 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3314,7 +3314,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. @@ -3535,6 +3537,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::(), 3); + Ok(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 49725983a391b..c28b6f6dc8872 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -1872,7 +1872,11 @@ fn build_predicate_expression( if !matches!(case.data_type(schema), Ok(DataType::Boolean)) { return unhandled_hook.handle(expr); } - // A missing `ELSE` yields NULL, which never matches, so it adds nothing. + // Without an `ELSE`, unmatched rows are UNKNOWN: not a match, but not + // FALSE either, which full-match inference must be able to tell apart. + if case.else_expr().is_none() { + return unhandled_hook.handle(expr); + } return case .when_then_expr() .iter() diff --git a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt index 2bdb99acfe3d3..1dc66f398d40d 100644 --- a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt +++ b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt @@ -1052,3 +1052,50 @@ set datafusion.execution.parquet.reorder_filters = false; statement ok DROP TABLE dict_filter_bug; + +# Regression: a `CASE` with no `ELSE` leaves unmatched rows UNKNOWN, not FALSE. +# Treating them as FALSE lets the inverse predicate call a row group fully +# matched and skip the row filter. Required columns admit no `IS NULL` rescue. +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +COPY (SELECT column1 AS a, column2 AS b FROM (VALUES (2, 1), (1, 2))) +TO 'test_files/scratch/parquet_filter_pushdown/case_without_else.parquet' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' '1'); + +statement ok +CREATE EXTERNAL TABLE case_without_else (a BIGINT NOT NULL, b BIGINT NOT NULL) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/case_without_else.parquet'; + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN false END) AND b > 0 ORDER BY a; +---- +1 2 + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN false END) AND b > 0 LIMIT 1; +---- +1 2 + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN false ELSE NULL END) AND b > 0 ORDER BY a; +---- +1 2 + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN NULL ELSE false END) AND b > 0 ORDER BY a; +---- +2 1 + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +statement ok +DROP TABLE case_without_else; From 8b8b27fba634876d1ee8520cf2abbfd1fcfbb15b Mon Sep 17 00:00:00 2001 From: Marko Grujic Date: Thu, 24 Sep 2026 07:07:46 +0200 Subject: [PATCH 3/6] Update datafusion/physical-plan/src/joins/hash_join/exec.rs Co-authored-by: Jay Zhan --- .../physical-plan/src/joins/hash_join/exec.rs | 86 ++++++++++--------- 1 file changed, 47 insertions(+), 39 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 3a9267160507f..d784bd78a73e2 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3355,11 +3355,15 @@ async fn collect_left_input( .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()); - } + // 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. + let pruning_bitmap = pruning_bitmap.filter(|bitmap| { + let ok = reservation.try_grow(bitmap.size()).is_ok(); + if ok { + metrics.build_mem_used.add(bitmap.size()); + } + ok + }); PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap) } }; @@ -3537,50 +3541,54 @@ mod tests { Ok(()) } - #[tokio::test] - async fn no_pruning_state_without_dynamic_filters() -> Result<()> { + /// 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. + async fn build_mem_used(limit: usize, dynamic_filters: bool) -> 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 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 = 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, - )?; + let build = batch((0..200).map(|i| i * 10_000).collect())?; + let probe = batch(vec![0, 500_000, 1_500_000])?; + let join = if dynamic_filters { + hash_join_with_dynamic_filter(build, probe, on, JoinType::Inner)?.0 + } else { + join( + build, + probe, + 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) + .with_memory_limit(limit, 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 task_ctx = Arc::new(TaskContext::default().with_runtime(runtime)); let batches = common::collect(join.execute(0, task_ctx)?).await?; assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 3); + Ok(join + .metrics() + .unwrap() + .sum_by_name("build_mem_used") + .unwrap() + .as_usize()) + } + + /// The bitmap is pruning-only: never built when no dynamic filter will read + /// it, and dropped rather than fatal when the pool cannot fit it. + #[tokio::test] + async fn pruning_bitmap_is_optional() -> Result<()> { + assert!(build_mem_used(1_000_000, false).await? < 100_000); + assert!(build_mem_used(1_000_000, true).await? > 100_000); + build_mem_used(100_000, true).await?; Ok(()) } From f1c7b038d24e56c7c183a751f3d255c789c03d07 Mon Sep 17 00:00:00 2001 From: Marko Grujic Date: Thu, 24 Sep 2026 16:28:17 +0200 Subject: [PATCH 4/6] See through a cast in the hash-lookup pruning expr Co-authored-by: Jay Zhan Co-Authored-By: Claude Sonnet 5 --- Cargo.lock | 1 + .../joins/hash_join/exec/prepared/tests.rs | 4 +- datafusion/pruning/Cargo.toml | 1 + datafusion/pruning/src/pruning_predicate.rs | 68 ++++++++++++++++++- .../test_files/dynamic_row_group_pruning.slt | 5 +- 5 files changed, 74 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0ea11794b16c6..3b5ed185cbd8d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2642,6 +2642,7 @@ dependencies = [ "insta", "itertools 0.15.0", "log", + "rstest", ] [[package]] diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs index e83f953a4c9f0..9e50b33deed63 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs @@ -795,7 +795,7 @@ async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { Arc::new(config), ) .await?; - assert!(matches!(map.build.membership, PushdownStrategy::Map(_))); + assert!(matches!(map.build.membership, PushdownStrategy::Map(_, _))); assert_eq!(map.reserved_bytes(), bytes); assert_eq!(pool.reserved(), 2 * bytes); drop(map); @@ -888,7 +888,7 @@ async fn prepared_byte_keys_use_hash_membership() -> Result<()> { assert!(prepared.build.bounds.is_none()); assert!(matches!( &prepared.build.membership, - PushdownStrategy::Map(_) + PushdownStrategy::Map(_, _) )); let task = with_probe_filter(base)? .builder() diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index a914a0a079bcd..48b1ed009b28c 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -31,6 +31,7 @@ datafusion-expr = { workspace = true } datafusion-functions-nested = { workspace = true } insta = { workspace = true } itertools = { workspace = true } +rstest = { workspace = true } [[bench]] harness = false diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index c28b6f6dc8872..4021f25d3ecd9 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -1540,8 +1540,14 @@ fn min_max_non_null_exprs( let max = required_columns .max_column_expr(column, column_expr, field) .ok()?; + // A cast on `column_expr` (widening it towards min/max's target type) + // does not affect nullability, so check it on the bare column beneath. + let bare_expr = match column_expr.downcast_ref::() { + Some(cast) => cast.expr(), + None => column_expr, + }; let non_null = - build_is_null_column_expr(column_expr, schema, required_columns, true)?; + build_is_null_column_expr(bare_expr, schema, required_columns, true)?; Some((min, max, non_null)) })(); if statistics.is_none() { @@ -1564,7 +1570,17 @@ fn build_hash_lookup_pruning_expr( let [column_expr] = on_columns[..] else { return None; }; - let column = column_expr.downcast_ref::()?; + // The schema adapter often wraps the column in a same-type or widening cast + // (a nullable table column over a REQUIRED file column, or a narrower file + // type). `stat_column_expr` rewrites only the inner `Column`, keeping it. + let column = match column_expr.downcast_ref::() { + Some(cast) => { + let column = cast.expr().downcast_ref::()?; + let from = schema.fields().get(column.index())?.data_type(); + cast.is_bigger_cast(from).then_some(column)? + } + None => column_expr.downcast_ref::()?, + }; let field = schema.fields().get(column.index())?; if field.name() != column.name() { return None; @@ -2492,6 +2508,7 @@ mod tests { use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; use datafusion_physical_plan::joins::{Map, SeededRandomState}; use itertools::Itertools; + use rstest::rstest; #[derive(Debug, Default)] /// Mock statistic provider for tests @@ -7476,6 +7493,53 @@ mod tests { assert_eq!(result, vec![false, true, false]); } + #[rstest] + #[case::widening_i32_to_i64(DataType::Int32, DataType::Int64, true)] + #[case::widening_u32_to_i64(DataType::UInt32, DataType::Int64, true)] + #[case::widening_u16_to_i64(DataType::UInt16, DataType::Int64, true)] + #[case::same_type_i64_to_i64(DataType::Int64, DataType::Int64, true)] + #[case::narrowing_i64_to_i32(DataType::Int64, DataType::Int32, false)] + #[case::narrowing_u32_to_i16(DataType::UInt32, DataType::Int16, false)] + fn test_hash_lookup_pruning_through_cast( + #[case] file_type: DataType, + #[case] cast_type: DataType, + #[case] pruned: bool, + ) { + let schema = + Arc::new(Schema::new(vec![Field::new("b", file_type.clone(), true)])); + let column: Arc = Arc::new(phys_expr::Column::new("b", 0)); + let cast_column: Arc = + Arc::new(phys_expr::CastExpr::new(column, cast_type, None)); + let lookup = hash_lookup(&cast_column, &[10, 20, 30, 100_000]); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(lookup) + .unwrap(); + + let to_file_type = |values: Vec>| -> ArrayRef { + let values: ArrayRef = Arc::new(values.into_iter().collect::()); + arrow::compute::cast(&values, &file_type).unwrap() + }; + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new() + .with_min(to_file_type(vec![Some(5000), Some(15), Some(50000)])) + .with_max(to_file_type(vec![Some(8000), Some(25), Some(60000)])), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 + // holds 20. A working cast excludes the first and third; a declined one + // falls back to keeping every container. + let expected = if pruned { + vec![false, true, false] + } else { + vec![true, true, true] + }; + assert_eq!(result, expected); + } + #[test] fn test_partition_routed_hash_lookup_pruning() { let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); diff --git a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt index fec01855094f7..753bf6fdfd263 100644 --- a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt +++ b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt @@ -269,7 +269,7 @@ SELECT WHEN i < 200 THEN 2000 + (i - 100) WHEN i < 300 THEN 3000 + (i - 200) ELSE (i - 300) - END AS BIGINT) AS b + END AS INT) AS b FROM generate_series(0, 399) AS t(i); statement ok @@ -285,6 +285,9 @@ OPTIONS ( statement ok drop table rgsel_src; +# `b` is stored as INT32 but declared BIGINT, so the schema adapter wraps every +# reference in a widening `CAST(b AS BIGINT)`. The hash-lookup pruning expr must +# see through that cast rather than silently losing the bitmap below. statement ok CREATE EXTERNAL TABLE rgsel (a BIGINT NOT NULL, b BIGINT NOT NULL) STORED AS PARQUET From 324f67f7ef7843bf256a5f444a383199b66981b9 Mon Sep 17 00:00:00 2001 From: Marko Grujic Date: Tue, 29 Sep 2026 07:43:06 +0200 Subject: [PATCH 5/6] Always compute pushdown membership for prepared builds Skipping unused pruning state made membership depend on `should_compute_dynamic_filters`, but a prepared build doesn't know whether a future consuming join will attach a dynamic filter, so it must always compute it. Also updates a pre-existing test whose byte accounting assumed a bitmap is never added on top of the map. Co-Authored-By: Claude Sonnet 5 --- .../physical-plan/src/joins/hash_join/exec.rs | 7 ++++--- .../src/joins/hash_join/exec/prepared/tests.rs | 13 +++++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index d784bd78a73e2..b69345f867d45 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3314,9 +3314,10 @@ async fn collect_left_input( let map = Arc::new(join_hash_map); - // 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 { + // For an ordinary build, nothing reads the strategy unless this join's own + // dynamic filter accumulator exists. A prepared build is shared with joins + // decided later, so it must always be ready to serve one. + let membership = if num_rows == 0 || (!prepared && !should_compute_dynamic_filters) { PushdownStrategy::Empty } else { // If the build side is small enough we can use IN list pushdown. diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs index 9e50b33deed63..48281e5f481b1 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs @@ -785,7 +785,9 @@ async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { assert_eq!(counter.count_batch(&membership_batch), 0); drop(membership_batch); - // Disabling IN-list membership changes no retained key buffers. + // Disabling IN-list membership shares the same key buffers; a + // pruning bitmap, when the sparse two-key range earns one, adds + // its own bytes on top. let mut config = ConfigOptions::default(); config.optimizer.hash_join_inlist_pushdown_max_size = 0; let map = base @@ -795,9 +797,12 @@ async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { Arc::new(config), ) .await?; - assert!(matches!(map.build.membership, PushdownStrategy::Map(_, _))); - assert_eq!(map.reserved_bytes(), bytes); - assert_eq!(pool.reserved(), 2 * bytes); + let PushdownStrategy::Map(_, bitmap) = &map.build.membership else { + panic!("expected hash-map membership"); + }; + let bitmap_bytes = bitmap.as_ref().map_or(0, |bitmap| bitmap.size()); + assert_eq!(map.reserved_bytes(), bytes + bitmap_bytes); + assert_eq!(pool.reserved(), 2 * bytes + bitmap_bytes); drop(map); let lease = Arc::clone(&prepared); drop(prepared); From 69867be6bdb8657647bb97716db0f6e88c7dcc44 Mon Sep 17 00:00:00 2001 From: Marko Grujic Date: Wed, 30 Sep 2026 16:32:45 +0200 Subject: [PATCH 6/6] Make pruning_bitmap_is_optional robust to force_hash_collisions Under `force_hash_collisions` every key hashes to the same bucket, so Map::num_of_distinct_key() collapses to 1 and the IN-list threshold check no longer reflects the test's real 200 distinct keys, sending it down the IN-list branch instead of the Map+bitmap branch under test. Force the threshold off directly instead of relying on distinct-key counting to cross it. Co-Authored-By: Claude Sonnet 5 --- .../physical-plan/src/joins/hash_join/exec.rs | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index b69345f867d45..004d7ed448dce 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -3544,6 +3544,11 @@ mod tests { /// 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 { let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); let batch = |keys: Vec| { @@ -3572,7 +3577,16 @@ mod tests { let runtime = RuntimeEnvBuilder::new() .with_memory_limit(limit, 1.0) .build_arc()?; - let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime)); + 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); Ok(join