diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 9858a4c06cd4a..11309d911a4a3 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, @@ -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::(); - 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::(), 3); + Ok(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, 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 82863c080259f..1b5674d63cbb7 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 { @@ -374,12 +385,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 { @@ -425,6 +437,7 @@ impl PhysicalExpr for HashTableLookupExpr { random_state: _, map: _, description: _, + pruning_bitmap: _, } = self; // HashTableLookupExpr holds a runtime Arc (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::().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,6 +803,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -766,6 +811,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_eq!(expr1, expr2); @@ -784,6 +830,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -791,6 +838,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -807,6 +855,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_one".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -814,6 +863,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_two".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -833,6 +883,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map1, "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -840,6 +891,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map2, "lookup".to_string(), + None, ); // Different Arc pointers means not equal (uses Arc::ptr_eq) @@ -857,6 +909,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -864,6 +917,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 10e5793e0ff32..775cb1ac11b4d 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -54,6 +54,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..c28b6f6dc8872 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,45 @@ 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); + } + // 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() + .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 +2488,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 +7424,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; 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;