From 8964e718b03324a895da3d085e5df624805cbe30 Mon Sep 17 00:00:00 2001 From: Devan Date: Mon, 12 Oct 2026 17:41:11 -0500 Subject: [PATCH 1/7] feat(perf): compute routing and lookup hashes for case once --- .../joins/hash_join/partitioned_hash_eval.rs | 271 +++++++++++++++++- .../src/joins/hash_join/shared_bounds.rs | 102 ++++++- 2 files changed, 367 insertions(+), 6 deletions(-) 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 921a38970562e..d002184967a26 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 @@ -20,15 +20,13 @@ use std::{fmt::Display, hash::Hash, sync::Arc}; use arrow::{ - array::{ArrayRef, UInt64Array}, + array::{ArrayRef, BooleanArray, UInt64Array}, datatypes::{DataType, Schema}, record_batch::RecordBatch, }; -use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::{create_hashes, with_hashes}; -#[cfg(feature = "proto")] -use datafusion_common::internal_err; +use datafusion_common::{Result, internal_err}; use datafusion_expr::ColumnarValue; use datafusion_expr_common::dyn_eq::DynHash; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; @@ -458,6 +456,209 @@ impl PhysicalExpr for HashTableLookupExpr { } } +/// Build-side lookup for one partition of a `Partitioned` hash join. +pub(crate) enum PartitionLookup { + /// The partition's hash table; probe rows routed here are looked up in it. + Map(Arc), + /// The partition is known empty (`false`) or its contents are unknown (`true`). + Const(bool), +} + +/// Routes each probe row to the hash table of the build partition that would +/// contain it and looks the row up there. +/// +/// This replaces `CASE hash(keys) % n WHEN i THEN lookup_i ... END` for +/// partitioned joins. `CaseExpr` evaluates each branch by filtering the batch +/// down to that branch's rows and merging the partial results, which copies +/// every column of the batch once per branch. This expression computes the +/// routing and lookup hashes once, buckets row indices by partition and only +/// gathers the 8-byte hashes, so the batch columns are never copied. +pub(crate) struct PartitionedLookupExpr { + on_columns: Vec, + routing_state: SeededRandomState, + lookup_state: SeededRandomState, + partitions: Vec, +} + +impl PartitionedLookupExpr { + pub(crate) fn new( + on_columns: Vec, + routing_state: SeededRandomState, + lookup_state: SeededRandomState, + partitions: Vec, + ) -> Self { + Self { + on_columns, + routing_state, + lookup_state, + partitions, + } + } +} + +impl std::fmt::Debug for PartitionedLookupExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let cols = self + .on_columns + .iter() + .map(|e| e.to_string()) + .collect::>() + .join(", "); + write!( + f, + "partitioned_hash_lookup({cols}, {} partitions)", + self.partitions.len() + ) + } +} + +impl Display for PartitionedLookupExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "partitioned_hash_lookup({})", self.partitions.len()) + } +} + +impl Hash for PartitionedLookupExpr { + fn hash(&self, state: &mut H) { + self.on_columns.dyn_hash(state); + self.routing_state.seed().hash(state); + self.lookup_state.seed().hash(state); + for partition in &self.partitions { + match partition { + PartitionLookup::Map(map) => Arc::as_ptr(map).hash(state), + PartitionLookup::Const(value) => value.hash(state), + } + } + } +} + +impl PartialEq for PartitionedLookupExpr { + fn eq(&self, other: &Self) -> bool { + self.on_columns == other.on_columns + && self.routing_state.seed() == other.routing_state.seed() + && self.lookup_state.seed() == other.lookup_state.seed() + && self.partitions.len() == other.partitions.len() + && self + .partitions + .iter() + .zip(&other.partitions) + .all(|(a, b)| match (a, b) { + (PartitionLookup::Map(a), PartitionLookup::Map(b)) => { + Arc::ptr_eq(a, b) + } + (PartitionLookup::Const(a), PartitionLookup::Const(b)) => a == b, + _ => false, + }) + } +} + +impl Eq for PartitionedLookupExpr {} + +impl PhysicalExpr for PartitionedLookupExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let keys = evaluate_columns(&self.on_columns, batch)?; + let num_rows = batch.num_rows(); + let num_partitions = self.partitions.len() as u64; + + // Same routing as RepartitionExec: hash(keys) % partitions. + let mut routing_hashes = vec![0u64; num_rows]; + create_hashes( + &keys, + self.routing_state.random_state(), + &mut routing_hashes, + )?; + let mut rows_by_partition: Vec> = + vec![Vec::new(); self.partitions.len()]; + for (row, hash) in routing_hashes.iter().enumerate() { + rows_by_partition[(hash % num_partitions) as usize].push(row); + } + + with_hashes(&keys, self.lookup_state.random_state(), |lookup_hashes| { + let mut result = vec![false; num_rows]; + for (partition, rows) in self.partitions.iter().zip(&rows_by_partition) { + match partition { + PartitionLookup::Const(value) => { + for &row in rows { + result[row] = *value; + } + } + PartitionLookup::Map(map) => { + let Map::HashMap(map) = map.as_ref() else { + return internal_err!( + "PartitionedLookupExpr only supports hash map partitions" + ); + }; + let hashes: Vec = + rows.iter().map(|&row| lookup_hashes[row]).collect(); + let hits = map.contain_hashes(&hashes); + for (i, &row) in rows.iter().enumerate() { + result[row] = hits.value(i); + } + } + } + } + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(result)))) + }) + } + + fn children(&self) -> Vec<&Arc> { + self.on_columns.iter().collect() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let partitions = self + .partitions + .iter() + .map(|p| match p { + PartitionLookup::Map(map) => PartitionLookup::Map(Arc::clone(map)), + PartitionLookup::Const(value) => PartitionLookup::Const(*value), + }) + .collect(); + Ok(Arc::new(PartitionedLookupExpr { + on_columns: children, + routing_state: self.routing_state.clone(), + lookup_state: self.lookup_state.clone(), + partitions, + })) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{self}") + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use datafusion_proto_models::protobuf::physical_expr_node::ExprType; + + // Like HashTableLookupExpr: the hash tables cannot be serialized, and + // dynamic filters are an optimization only, so emit lit(true). + let value = datafusion_proto_common::ScalarValue { + value: Some(datafusion_proto_common::scalar_value::Value::BoolValue( + true, + )), + }; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(ExprType::Literal(value)), + })) + } +} + fn evaluate_columns( columns: &[PhysicalExprRef], batch: &RecordBatch, @@ -472,7 +673,7 @@ fn evaluate_columns( #[cfg(test)] mod tests { use super::*; - use crate::joins::join_hash_map::JoinHashMapU32; + use crate::joins::join_hash_map::{JoinHashMapType, JoinHashMapU32}; use datafusion_common::ScalarValue; use datafusion_physical_expr::expressions::Column; use std::collections::hash_map::DefaultHasher; @@ -514,6 +715,66 @@ mod tests { assert!(Arc::ptr_eq(&bitmap, derived.pruning_bitmap().unwrap())); } + /// Rows must be looked up in the table of the partition RepartitionExec + /// would send them to, and constant partitions must apply to their rows only. + #[test] + fn partitioned_lookup_routes_rows_to_their_partition() { + let routing = SeededRandomState::with_seed(1); + let lookup = SeededRandomState::with_seed(2); + let num_partitions = 3u64; + let keys: ArrayRef = + Arc::new(arrow::array::Int32Array::from((0..300).collect::>())); + let mut routing_hashes = vec![0u64; 300]; + create_hashes( + &[Arc::clone(&keys)], + routing.random_state(), + &mut routing_hashes, + ) + .unwrap(); + let mut lookup_hashes = vec![0u64; 300]; + create_hashes( + &[Arc::clone(&keys)], + lookup.random_state(), + &mut lookup_hashes, + ) + .unwrap(); + + // Partition 0 and 1 hold the even keys routed to them; partition 2 is a constant. + let partitions = (0..num_partitions) + .map(|p| { + if p == 2 { + return PartitionLookup::Const(true); + } + let build: Vec = (0..300) + .filter(|k| k % 2 == 0 && routing_hashes[*k] % num_partitions == p) + .map(|k| lookup_hashes[k]) + .collect(); + let mut map = JoinHashMapU32::with_capacity(build.len()); + map.update_from_iter(Box::new(build.iter().enumerate()), 0); + PartitionLookup::Map(Arc::new(Map::HashMap(Box::new(map)))) + }) + .collect(); + let expr = PartitionedLookupExpr::new( + vec![Arc::new(Column::new("k", 0))], + routing, + lookup, + partitions, + ); + + let schema = Arc::new(Schema::new(vec![arrow::datatypes::Field::new( + "k", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new(schema, vec![keys]).unwrap(); + let result = expr.evaluate(&batch).unwrap().into_array(300).unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + for (k, routing_hash) in routing_hashes.iter().enumerate() { + let expected = routing_hash % num_partitions == 2 || k % 2 == 0; + assert_eq!(result.value(k), expected, "key {k}"); + } + } + #[test] fn test_hash_expr_eq_same() { let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); 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 7d81dc9162ca7..ef2547159b3cc 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -29,7 +29,8 @@ use crate::joins::PartitionMode; use crate::joins::hash_join::exec::{HASH_JOIN_SEED, NullAwareMode}; use crate::joins::hash_join::inlist_builder::build_struct_fields; use crate::joins::hash_join::partitioned_hash_eval::{ - HashExpr, HashTableLookupExpr, SeededRandomState, + HashExpr, HashTableLookupExpr, PartitionLookup, PartitionedLookupExpr, + SeededRandomState, }; use crate::joins::key_range_bitmap::KeyRangeBitmap; use crate::repartition::RangeExpr; @@ -85,6 +86,25 @@ impl PartitionBounds { pub(crate) fn get_column_bounds(&self, index: usize) -> Option<&ColumnBounds> { self.column_bounds.get(index) } + + /// Bounds covering both `self` and `other`, or `None` if they do not + /// describe the same columns. + fn union(self, other: &PartitionBounds) -> Option { + if self.column_bounds.len() != other.column_bounds.len() { + return None; + } + let column_bounds = self + .column_bounds + .into_iter() + .zip(&other.column_bounds) + .map(|(a, b)| { + let min = if b.min < a.min { b.min.clone() } else { a.min }; + let max = if b.max > a.max { b.max.clone() } else { a.max }; + ColumnBounds::new(min, max) + }) + .collect(); + Some(PartitionBounds::new(column_bounds)) + } } /// Creates a membership predicate for filter pushdown. @@ -668,6 +688,10 @@ impl SharedBuildAccumulator { let mut empty_partition_ids = Vec::new(); let mut has_canceled_unknown = false; let mut keys_have_null = false; + // Per-partition lookups for `PartitionedLookupExpr`; `None` once a + // partition needs a predicate that expression cannot route to. + let mut lookups = Some(Vec::with_capacity(partition_filters.capacity())); + let mut merged_bounds: Option = None; for (partition_id, partition) in partitions.into_iter().enumerate() { match partition { @@ -677,6 +701,9 @@ impl SharedBuildAccumulator { }) => { empty_partition_ids.push(partition_id); partition_filters.push(lit(false)); + if let Some(lookups) = &mut lookups { + lookups.push(PartitionLookup::Const(false)); + } } PartitionStatus::Reported(PartitionData { bounds, @@ -685,6 +712,20 @@ impl SharedBuildAccumulator { }) => { real_partition_ids.push(partition_id); keys_have_null |= partition_keys_have_null; + lookups = match (&pushdown, lookups) { + (PushdownStrategy::Map(map, _), Some(mut lookups)) + if matches!(map.as_ref(), Map::HashMap(_)) => + { + lookups.push(PartitionLookup::Map(Arc::clone(map))); + Some(lookups) + } + _ => None, + }; + merged_bounds = match (merged_bounds, real_partition_ids.len()) { + (None, 1) => Some(bounds.clone()), + (Some(merged), _) => merged.union(&bounds), + (None, _) => None, + }; let membership_expr = create_membership_predicate( &self.on_right, pushdown, @@ -700,6 +741,9 @@ impl SharedBuildAccumulator { PartitionStatus::CanceledUnknown => { has_canceled_unknown = true; partition_filters.push(lit(true)); + if let Some(lookups) = &mut lookups { + lookups.push(PartitionLookup::Const(true)); + } // A canceled partition's build content is unknown, so it // may hold a NULL key. keys_have_null = true; @@ -729,6 +773,22 @@ impl SharedBuildAccumulator { // Only one build partition contains rows, so its filter covers every // possible probe match without routing. Arc::clone(&partition_filters[real_partition_ids[0]]) + } else if let Some(lookups) = lookups.filter(|lookups| { + self.probe_range_partitioning.is_none() + && lookups.iter().any(|l| matches!(l, PartitionLookup::Map(_))) + }) { + // Every partition is a hash table, so route rows to the right table + // directly instead of through a CASE over per-partition filters. + let membership = Arc::new(PartitionedLookupExpr::new( + self.on_right.clone(), + self.repartition_random_state.clone(), + HASH_JOIN_SEED, + lookups, + )) as Arc; + let bounds = merged_bounds + .and_then(|bounds| create_bounds_predicate(&self.on_right, &bounds)); + combine_membership_and_bounds(Some(membership), bounds) + .expect("membership predicate is always present") } else { // Builds the shared sparse `CASE` for partition filter routing. // Without cancellation, omitted branches are known empty and safely fall @@ -919,6 +979,7 @@ pub(super) fn completed_partitions_for_test(acc: &SharedBuildAccumulator) -> usi mod tests { use super::*; + use crate::joins::join_hash_map::JoinHashMapU32; use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array}; use arrow::compute::SortOptions; use arrow::record_batch::RecordBatch; @@ -1135,6 +1196,45 @@ mod tests { assert_literal_bool(&expr, true); } + fn hash_map() -> PushdownStrategy { + PushdownStrategy::Map( + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(1)))), + None, + ) + } + + #[test] + fn partitioned_hash_map_partitions_route_without_case() { + let acc = make_partitioned_expr_accumulator_for_test(3); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(hash_map(), bounds(5, 10)), + reported(PushdownStrategy::Empty, no_bounds()), + reported(hash_map(), bounds(1, 7)), + ])) + .unwrap(); + + let expr = current_expr(&acc); + assert!(expr.downcast_ref::().is_none()); + assert_eq!( + expr.to_string(), + "probe_key@0 >= 1 AND probe_key@0 <= 10 AND partitioned_hash_lookup(3)" + ); + } + + #[test] + fn partitioned_in_list_partition_keeps_case() { + let acc = make_partitioned_expr_accumulator_for_test(2); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(hash_map(), no_bounds()), + reported(in_list(&[2]), no_bounds()), + ])) + .unwrap(); + + assert!(current_expr(&acc).downcast_ref::().is_some()); + } + #[test] fn partitioned_one_real_partition_with_rest_empty_skips_case() { let acc = make_partitioned_expr_accumulator_for_test(3); From c43a0077b9f3aed0cae4b15161fa053fdbf9f236 Mon Sep 17 00:00:00 2001 From: Devan Benz Date: Mon, 19 Oct 2026 18:05:29 -0500 Subject: [PATCH 2/7] feat(perf): Let's use HashTableLookupExpr instead so we get pruning --- .../joins/hash_join/partitioned_hash_eval.rs | 269 +----------------- .../src/joins/hash_join/shared_bounds.rs | 55 ++-- .../physical-plan/src/joins/join_hash_map.rs | 93 ++++-- .../src/joins/stream_join_utils.rs | 4 + 4 files changed, 110 insertions(+), 311 deletions(-) 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 d002184967a26..4f47e91342d6c 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 @@ -20,13 +20,13 @@ use std::{fmt::Display, hash::Hash, sync::Arc}; use arrow::{ - array::{ArrayRef, BooleanArray, UInt64Array}, + array::{ArrayRef, UInt64Array}, datatypes::{DataType, Schema}, record_batch::RecordBatch, }; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::{create_hashes, with_hashes}; -use datafusion_common::{Result, internal_err}; +use datafusion_common::Result; use datafusion_expr::ColumnarValue; use datafusion_expr_common::dyn_eq::DynHash; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; @@ -456,209 +456,6 @@ impl PhysicalExpr for HashTableLookupExpr { } } -/// Build-side lookup for one partition of a `Partitioned` hash join. -pub(crate) enum PartitionLookup { - /// The partition's hash table; probe rows routed here are looked up in it. - Map(Arc), - /// The partition is known empty (`false`) or its contents are unknown (`true`). - Const(bool), -} - -/// Routes each probe row to the hash table of the build partition that would -/// contain it and looks the row up there. -/// -/// This replaces `CASE hash(keys) % n WHEN i THEN lookup_i ... END` for -/// partitioned joins. `CaseExpr` evaluates each branch by filtering the batch -/// down to that branch's rows and merging the partial results, which copies -/// every column of the batch once per branch. This expression computes the -/// routing and lookup hashes once, buckets row indices by partition and only -/// gathers the 8-byte hashes, so the batch columns are never copied. -pub(crate) struct PartitionedLookupExpr { - on_columns: Vec, - routing_state: SeededRandomState, - lookup_state: SeededRandomState, - partitions: Vec, -} - -impl PartitionedLookupExpr { - pub(crate) fn new( - on_columns: Vec, - routing_state: SeededRandomState, - lookup_state: SeededRandomState, - partitions: Vec, - ) -> Self { - Self { - on_columns, - routing_state, - lookup_state, - partitions, - } - } -} - -impl std::fmt::Debug for PartitionedLookupExpr { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let cols = self - .on_columns - .iter() - .map(|e| e.to_string()) - .collect::>() - .join(", "); - write!( - f, - "partitioned_hash_lookup({cols}, {} partitions)", - self.partitions.len() - ) - } -} - -impl Display for PartitionedLookupExpr { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "partitioned_hash_lookup({})", self.partitions.len()) - } -} - -impl Hash for PartitionedLookupExpr { - fn hash(&self, state: &mut H) { - self.on_columns.dyn_hash(state); - self.routing_state.seed().hash(state); - self.lookup_state.seed().hash(state); - for partition in &self.partitions { - match partition { - PartitionLookup::Map(map) => Arc::as_ptr(map).hash(state), - PartitionLookup::Const(value) => value.hash(state), - } - } - } -} - -impl PartialEq for PartitionedLookupExpr { - fn eq(&self, other: &Self) -> bool { - self.on_columns == other.on_columns - && self.routing_state.seed() == other.routing_state.seed() - && self.lookup_state.seed() == other.lookup_state.seed() - && self.partitions.len() == other.partitions.len() - && self - .partitions - .iter() - .zip(&other.partitions) - .all(|(a, b)| match (a, b) { - (PartitionLookup::Map(a), PartitionLookup::Map(b)) => { - Arc::ptr_eq(a, b) - } - (PartitionLookup::Const(a), PartitionLookup::Const(b)) => a == b, - _ => false, - }) - } -} - -impl Eq for PartitionedLookupExpr {} - -impl PhysicalExpr for PartitionedLookupExpr { - fn data_type(&self, _input_schema: &Schema) -> Result { - Ok(DataType::Boolean) - } - - fn nullable(&self, _input_schema: &Schema) -> Result { - Ok(false) - } - - fn evaluate(&self, batch: &RecordBatch) -> Result { - let keys = evaluate_columns(&self.on_columns, batch)?; - let num_rows = batch.num_rows(); - let num_partitions = self.partitions.len() as u64; - - // Same routing as RepartitionExec: hash(keys) % partitions. - let mut routing_hashes = vec![0u64; num_rows]; - create_hashes( - &keys, - self.routing_state.random_state(), - &mut routing_hashes, - )?; - let mut rows_by_partition: Vec> = - vec![Vec::new(); self.partitions.len()]; - for (row, hash) in routing_hashes.iter().enumerate() { - rows_by_partition[(hash % num_partitions) as usize].push(row); - } - - with_hashes(&keys, self.lookup_state.random_state(), |lookup_hashes| { - let mut result = vec![false; num_rows]; - for (partition, rows) in self.partitions.iter().zip(&rows_by_partition) { - match partition { - PartitionLookup::Const(value) => { - for &row in rows { - result[row] = *value; - } - } - PartitionLookup::Map(map) => { - let Map::HashMap(map) = map.as_ref() else { - return internal_err!( - "PartitionedLookupExpr only supports hash map partitions" - ); - }; - let hashes: Vec = - rows.iter().map(|&row| lookup_hashes[row]).collect(); - let hits = map.contain_hashes(&hashes); - for (i, &row) in rows.iter().enumerate() { - result[row] = hits.value(i); - } - } - } - } - Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(result)))) - }) - } - - fn children(&self) -> Vec<&Arc> { - self.on_columns.iter().collect() - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - let partitions = self - .partitions - .iter() - .map(|p| match p { - PartitionLookup::Map(map) => PartitionLookup::Map(Arc::clone(map)), - PartitionLookup::Const(value) => PartitionLookup::Const(*value), - }) - .collect(); - Ok(Arc::new(PartitionedLookupExpr { - on_columns: children, - routing_state: self.routing_state.clone(), - lookup_state: self.lookup_state.clone(), - partitions, - })) - } - - fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{self}") - } - - #[cfg(feature = "proto")] - fn try_to_proto( - &self, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, - ) -> Result> { - use datafusion_proto_models::protobuf; - use datafusion_proto_models::protobuf::physical_expr_node::ExprType; - - // Like HashTableLookupExpr: the hash tables cannot be serialized, and - // dynamic filters are an optimization only, so emit lit(true). - let value = datafusion_proto_common::ScalarValue { - value: Some(datafusion_proto_common::scalar_value::Value::BoolValue( - true, - )), - }; - Ok(Some(protobuf::PhysicalExprNode { - expr_id: None, - expr_type: Some(ExprType::Literal(value)), - })) - } -} - fn evaluate_columns( columns: &[PhysicalExprRef], batch: &RecordBatch, @@ -673,7 +470,7 @@ fn evaluate_columns( #[cfg(test)] mod tests { use super::*; - use crate::joins::join_hash_map::{JoinHashMapType, JoinHashMapU32}; + use crate::joins::join_hash_map::JoinHashMapU32; use datafusion_common::ScalarValue; use datafusion_physical_expr::expressions::Column; use std::collections::hash_map::DefaultHasher; @@ -715,66 +512,6 @@ mod tests { assert!(Arc::ptr_eq(&bitmap, derived.pruning_bitmap().unwrap())); } - /// Rows must be looked up in the table of the partition RepartitionExec - /// would send them to, and constant partitions must apply to their rows only. - #[test] - fn partitioned_lookup_routes_rows_to_their_partition() { - let routing = SeededRandomState::with_seed(1); - let lookup = SeededRandomState::with_seed(2); - let num_partitions = 3u64; - let keys: ArrayRef = - Arc::new(arrow::array::Int32Array::from((0..300).collect::>())); - let mut routing_hashes = vec![0u64; 300]; - create_hashes( - &[Arc::clone(&keys)], - routing.random_state(), - &mut routing_hashes, - ) - .unwrap(); - let mut lookup_hashes = vec![0u64; 300]; - create_hashes( - &[Arc::clone(&keys)], - lookup.random_state(), - &mut lookup_hashes, - ) - .unwrap(); - - // Partition 0 and 1 hold the even keys routed to them; partition 2 is a constant. - let partitions = (0..num_partitions) - .map(|p| { - if p == 2 { - return PartitionLookup::Const(true); - } - let build: Vec = (0..300) - .filter(|k| k % 2 == 0 && routing_hashes[*k] % num_partitions == p) - .map(|k| lookup_hashes[k]) - .collect(); - let mut map = JoinHashMapU32::with_capacity(build.len()); - map.update_from_iter(Box::new(build.iter().enumerate()), 0); - PartitionLookup::Map(Arc::new(Map::HashMap(Box::new(map)))) - }) - .collect(); - let expr = PartitionedLookupExpr::new( - vec![Arc::new(Column::new("k", 0))], - routing, - lookup, - partitions, - ); - - let schema = Arc::new(Schema::new(vec![arrow::datatypes::Field::new( - "k", - DataType::Int32, - false, - )])); - let batch = RecordBatch::try_new(schema, vec![keys]).unwrap(); - let result = expr.evaluate(&batch).unwrap().into_array(300).unwrap(); - let result = result.as_any().downcast_ref::().unwrap(); - for (k, routing_hash) in routing_hashes.iter().enumerate() { - let expected = routing_hash % num_partitions == 2 || k % 2 == 0; - assert_eq!(result.value(k), expected, "key {k}"); - } - } - #[test] fn test_hash_expr_eq_same() { let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); 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 8cfe3a9aec45f..cc164498ef0d3 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -29,9 +29,9 @@ use crate::joins::PartitionMode; use crate::joins::hash_join::exec::{HASH_JOIN_SEED, NullAwareMode}; use crate::joins::hash_join::inlist_builder::build_struct_fields; use crate::joins::hash_join::partitioned_hash_eval::{ - HashExpr, HashTableLookupExpr, PartitionLookup, PartitionedLookupExpr, - SeededRandomState, + HashExpr, HashTableLookupExpr, SeededRandomState, }; +use crate::joins::join_hash_map::union_hash_maps; use crate::joins::key_range_bitmap::KeyRangeBitmap; use crate::repartition::RangeExpr; use arrow::array::ArrayRef; @@ -688,9 +688,9 @@ impl SharedBuildAccumulator { let mut empty_partition_ids = Vec::new(); let mut has_canceled_unknown = false; let mut keys_have_null = false; - // Per-partition lookups for `PartitionedLookupExpr`; `None` once a - // partition needs a predicate that expression cannot route to. - let mut lookups = Some(Vec::with_capacity(partition_filters.capacity())); + // Every non-empty partition's hash table, merged into one membership + // lookup below; `None` once a partition's build side cannot be merged. + let mut hash_maps = Some(Vec::with_capacity(partitions.len())); let mut merged_bounds: Option = None; for (partition_id, partition) in partitions.into_iter().enumerate() { @@ -701,9 +701,6 @@ impl SharedBuildAccumulator { }) => { empty_partition_ids.push(partition_id); partition_filters.push(lit(false)); - if let Some(lookups) = &mut lookups { - lookups.push(PartitionLookup::Const(false)); - } } PartitionStatus::Reported(PartitionData { bounds, @@ -712,12 +709,12 @@ impl SharedBuildAccumulator { }) => { real_partition_ids.push(partition_id); keys_have_null |= partition_keys_have_null; - lookups = match (&pushdown, lookups) { - (PushdownStrategy::Map(map, _), Some(mut lookups)) + hash_maps = match (&pushdown, hash_maps) { + (PushdownStrategy::Map(map, _), Some(mut hash_maps)) if matches!(map.as_ref(), Map::HashMap(_)) => { - lookups.push(PartitionLookup::Map(Arc::clone(map))); - Some(lookups) + hash_maps.push(Arc::clone(map)); + Some(hash_maps) } _ => None, }; @@ -726,6 +723,7 @@ impl SharedBuildAccumulator { (Some(merged), _) => merged.union(&bounds), (None, _) => None, }; + let membership_expr = create_membership_predicate( &self.on_right, pushdown, @@ -741,9 +739,8 @@ impl SharedBuildAccumulator { PartitionStatus::CanceledUnknown => { has_canceled_unknown = true; partition_filters.push(lit(true)); - if let Some(lookups) = &mut lookups { - lookups.push(PartitionLookup::Const(true)); - } + // Rows routed here must pass, which a merged lookup cannot express. + hash_maps = None; // A canceled partition's build content is unknown, so it // may hold a NULL key. keys_have_null = true; @@ -773,17 +770,25 @@ impl SharedBuildAccumulator { // Only one build partition contains rows, so its filter covers every // possible probe match without routing. Arc::clone(&partition_filters[real_partition_ids[0]]) - } else if let Some(lookups) = lookups.filter(|lookups| { - self.probe_range_partitioning.is_none() - && lookups.iter().any(|l| matches!(l, PartitionLookup::Map(_))) - }) { - // Every partition is a hash table, so route rows to the right table - // directly instead of through a CASE over per-partition filters. - let membership = Arc::new(PartitionedLookupExpr::new( + } else if let Some(hash_maps) = hash_maps { + // Every non-empty partition is a hash table. A key is only ever built + // into the partition it is routed to, so membership in the union of + // all tables equals membership in the routed table, and needs neither + // a routing hash nor a CASE over per-partition filters. + let tables = hash_maps + .iter() + .filter_map(|map| match map.as_ref() { + Map::HashMap(table) => Some(table.as_ref()), + Map::ArrayMap(_) => None, + }) + .collect::>(); + let union = Arc::new(Map::HashMap(union_hash_maps(&tables))); + let membership = Arc::new(HashTableLookupExpr::new( self.on_right.clone(), - self.repartition_random_state.clone(), HASH_JOIN_SEED, - lookups, + union, + String::from("hash_lookup"), + None, )) as Arc; let bounds = merged_bounds .and_then(|bounds| create_bounds_predicate(&self.on_right, &bounds)); @@ -1218,7 +1223,7 @@ mod tests { assert!(expr.downcast_ref::().is_none()); assert_eq!( expr.to_string(), - "probe_key@0 >= 1 AND probe_key@0 <= 10 AND partitioned_hash_lookup(3)" + "probe_key@0 >= 1 AND probe_key@0 <= 10 AND hash_lookup" ); } diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index 3f58756fc38f7..5ebea1d96eb11 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -144,6 +144,9 @@ pub trait JoinHashMapType: Send + Sync { /// Returns the number of entries in the join hash map. fn len(&self) -> usize; + + /// Returns every distinct hash value stored in the map. + fn hash_values(&self) -> Box + '_>; } pub struct JoinHashMapU32 { @@ -230,6 +233,10 @@ impl JoinHashMapType for JoinHashMapU32 { fn len(&self) -> usize { self.map.len() } + + fn hash_values(&self) -> Box + '_> { + Box::new(self.map.iter().map(|&(hash, _)| hash)) + } } pub struct JoinHashMapU64 { @@ -316,6 +323,10 @@ impl JoinHashMapType for JoinHashMapU64 { fn len(&self) -> usize { self.map.len() } + + fn hash_values(&self) -> Box + '_> { + Box::new(self.map.iter().map(|&(hash, _)| hash)) + } } use crate::joins::MapOffset; @@ -335,29 +346,71 @@ where { let mut chained = false; for (row, &hash_value) in iter { - let entry = map.entry( - hash_value, - |&(hash, _)| hash_value == hash, - |&(hash, _)| hash, - ); + chained |= insert_hash(map, next, row, hash_value, deleted_offset); + } + chained +} - match entry { - Occupied(mut occupied_entry) => { - // Already exists: add index to next array - let (_, index) = occupied_entry.get_mut(); - let prev_index = *index; - // Store new value inside hashmap - *index = T::try_from(row + 1).unwrap(); - // Update chained Vec at `row` with previous value - next[row - deleted_offset] = prev_index; - chained = true; - } - Vacant(vacant_entry) => { - vacant_entry.insert((hash_value, T::try_from(row + 1).unwrap())); - } +/// Inserts `row` under `hash_value`, chaining it onto any row already stored +/// there. Returns `true` if it was chained. +#[inline] +fn insert_hash( + map: &mut HashTable<(u64, T)>, + next: &mut [T], + row: usize, + hash_value: u64, + deleted_offset: usize, +) -> bool +where + T: Copy + TryFrom, + >::Error: Debug, +{ + let entry = map.entry( + hash_value, + |&(hash, _)| hash_value == hash, + |&(hash, _)| hash, + ); + + match entry { + Occupied(mut occupied_entry) => { + // Already exists: add index to next array + let (_, index) = occupied_entry.get_mut(); + let prev_index = *index; + // Store new value inside hashmap + *index = T::try_from(row + 1).unwrap(); + // Update chained Vec at `row` with previous value + next[row - deleted_offset] = prev_index; + true + } + Vacant(vacant_entry) => { + vacant_entry.insert((hash_value, T::try_from(row + 1).unwrap())); + false } } - chained +} + +/// Builds one map holding the union of `maps`' hash values, so a single +/// [`JoinHashMapType::contain_hashes`] call checks membership in all of them. +/// +/// Its row indices number the merged hash values rather than rows of any build +/// batch, so it is only meaningful for membership checks, never to produce +/// join matches. +pub(crate) fn union_hash_maps(maps: &[&dyn JoinHashMapType]) -> Box { + let len = maps.iter().map(|map| map.len()).sum(); + let hashes = maps.iter().flat_map(|map| map.hash_values()); + if len > u32::MAX as usize { + let mut union = JoinHashMapU64::with_capacity(len); + for (row, hash) in hashes.enumerate() { + union.chained |= insert_hash(&mut union.map, &mut union.next, row, hash, 0); + } + Box::new(union) + } else { + let mut union = JoinHashMapU32::with_capacity(len); + for (row, hash) in hashes.enumerate() { + union.chained |= insert_hash(&mut union.map, &mut union.next, row, hash, 0); + } + Box::new(union) + } } pub fn get_matched_indices<'a, T>( diff --git a/datafusion/physical-plan/src/joins/stream_join_utils.rs b/datafusion/physical-plan/src/joins/stream_join_utils.rs index f74c3671d0557..be1c54aa63d42 100644 --- a/datafusion/physical-plan/src/joins/stream_join_utils.rs +++ b/datafusion/physical-plan/src/joins/stream_join_utils.rs @@ -117,6 +117,10 @@ impl JoinHashMapType for PruningJoinHashMap { fn len(&self) -> usize { self.map.len() } + + fn hash_values(&self) -> Box + '_> { + Box::new(self.map.iter().map(|&(hash, _)| hash)) + } } /// The `PruningJoinHashMap` is similar to a regular `JoinHashMap`, but with From b0505a05e630908153ae0f304bf937032fceb0cf Mon Sep 17 00:00:00 2001 From: Devan Benz Date: Mon, 19 Oct 2026 18:27:42 -0500 Subject: [PATCH 3/7] Merge per-partition pruning bitmaps --- .../joins/hash_join/partitioned_hash_eval.rs | 10 +- .../src/joins/hash_join/shared_bounds.rs | 45 ++++++-- .../src/joins/key_range_bitmap.rs | 109 ++++++++++++++++++ 3 files changed, 154 insertions(+), 10 deletions(-) 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 4f47e91342d6c..7d184b21a0944 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 @@ -24,9 +24,9 @@ use arrow::{ datatypes::{DataType, Schema}, record_batch::RecordBatch, }; +use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::{create_hashes, with_hashes}; -use datafusion_common::Result; use datafusion_expr::ColumnarValue; use datafusion_expr_common::dyn_eq::DynHash; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; @@ -251,7 +251,9 @@ impl HashExpr { let Some(protobuf::physical_expr_node::ExprType::HashExpr(hash_expr)) = &node.expr_type else { - return internal_err!("PhysicalExprNode is not a HashExpr"); + return datafusion_common::internal_err!( + "PhysicalExprNode is not a HashExpr" + ); }; // Destructure exhaustively (no `..`) so that a newly added proto field // is a compile error here instead of being silently ignored. @@ -653,7 +655,9 @@ mod tests { Some(protobuf::physical_expr_node::ExprType::Column(_)) => { Column::try_from_proto(node, &ctx) } - _ => internal_err!("test decoder cannot decode {node:?}"), + _ => datafusion_common::internal_err!( + "test decoder cannot decode {node:?}" + ), } } } 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 cc164498ef0d3..1581062cbad67 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -32,7 +32,7 @@ use crate::joins::hash_join::partitioned_hash_eval::{ HashExpr, HashTableLookupExpr, SeededRandomState, }; use crate::joins::join_hash_map::union_hash_maps; -use crate::joins::key_range_bitmap::KeyRangeBitmap; +use crate::joins::key_range_bitmap::{KeyRangeBitmap, KeyRangePart}; use crate::repartition::RangeExpr; use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, Schema}; @@ -688,8 +688,9 @@ impl SharedBuildAccumulator { let mut empty_partition_ids = Vec::new(); let mut has_canceled_unknown = false; let mut keys_have_null = false; - // Every non-empty partition's hash table, merged into one membership - // lookup below; `None` once a partition's build side cannot be merged. + // Every non-empty partition's hash table, pruning bitmap, and bounds, + // merged into one membership lookup below; `None` once a partition's + // build side cannot be merged. let mut hash_maps = Some(Vec::with_capacity(partitions.len())); let mut merged_bounds: Option = None; @@ -710,10 +711,14 @@ impl SharedBuildAccumulator { real_partition_ids.push(partition_id); keys_have_null |= partition_keys_have_null; hash_maps = match (&pushdown, hash_maps) { - (PushdownStrategy::Map(map, _), Some(mut hash_maps)) + (PushdownStrategy::Map(map, bitmap), Some(mut hash_maps)) if matches!(map.as_ref(), Map::HashMap(_)) => { - hash_maps.push(Arc::clone(map)); + hash_maps.push(( + Arc::clone(map), + bitmap.clone(), + bounds.clone(), + )); Some(hash_maps) } _ => None, @@ -777,18 +782,44 @@ impl SharedBuildAccumulator { // a routing hash nor a CASE over per-partition filters. let tables = hash_maps .iter() - .filter_map(|map| match map.as_ref() { + .filter_map(|(map, _, _)| match map.as_ref() { Map::HashMap(table) => Some(table.as_ref()), Map::ArrayMap(_) => None, }) .collect::>(); let union = Arc::new(Map::HashMap(union_hash_maps(&tables))); + // Container pruning reads one bitmap, so merge the partitions' bitmaps. + // A partition without one may hold a key anywhere in its bounds. Not + // worth building when no partition has a bitmap to contribute gaps. + let pruning_bitmap = merged_bounds + .as_ref() + .and_then(|merged| merged.get_column_bounds(0)) + .filter(|_| { + self.on_right.len() == 1 + && hash_maps.iter().any(|(_, bitmap, _)| bitmap.is_some()) + }) + .and_then(|merged| { + let parts = hash_maps + .iter() + .map(|(_, bitmap, bounds)| match bitmap { + Some(bitmap) => Some(KeyRangePart::Bitmap(bitmap)), + None => bounds.get_column_bounds(0).map(|bounds| { + KeyRangePart::Range { + min: &bounds.min, + max: &bounds.max, + } + }), + }) + .collect::>>()?; + KeyRangeBitmap::union(&merged.min, &merged.max, &parts) + }) + .map(Arc::new); let membership = Arc::new(HashTableLookupExpr::new( self.on_right.clone(), HASH_JOIN_SEED, union, String::from("hash_lookup"), - None, + pruning_bitmap, )) as Arc; let bounds = merged_bounds .and_then(|bounds| create_bounds_predicate(&self.on_right, &bounds)); diff --git a/datafusion/physical-plan/src/joins/key_range_bitmap.rs b/datafusion/physical-plan/src/joins/key_range_bitmap.rs index efdb5467c897e..bd87db415b558 100644 --- a/datafusion/physical-plan/src/joins/key_range_bitmap.rs +++ b/datafusion/physical-plan/src/joins/key_range_bitmap.rs @@ -28,6 +28,17 @@ 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; +/// One build partition's contribution to [`KeyRangeBitmap::union`]. +pub(crate) enum KeyRangePart<'a> { + /// The partition's own bitmap. + Bitmap(&'a KeyRangeBitmap), + /// The partition has no bitmap, so any key in `[min, max]` may be present. + Range { + min: &'a ScalarValue, + max: &'a ScalarValue, + }, +} + /// 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, @@ -101,6 +112,61 @@ impl KeyRangeBitmap { }) } + /// One bitmap over `[min, max]` marking every bucket that any of `parts` + /// may hold a key in. `[min, max]` must cover every part. + /// + /// Coarser than the parts' own bitmaps where its buckets are wider, but + /// never clears a bucket a part may hold a key in. `None` when the key + /// type has no `u64` ordering, or when nothing can be excluded. + pub(crate) fn union( + min: &ScalarValue, + max: &ScalarValue, + parts: &[KeyRangePart<'_>], + ) -> Option { + let offset = ArrayMap::key_to_u64(min)?; + let span = ArrayMap::key_to_u64(max)?.wrapping_sub(offset); + let buckets = span.saturating_add(1).min(MAX_BUCKETS).next_power_of_two(); + let shift = (0u32..63).find(|s| (span >> s) < buckets).unwrap_or(63); + + let mut words = MutableBuffer::new_null(buckets as usize); + let bytes = words.as_slice_mut(); + // Sets the buckets covering keys `lo..=hi`, given as offsets from `offset`. + let mut set_range = |lo: u64, hi: u64| { + for bucket in (lo >> shift)..=(hi.min(span) >> shift) { + bit_util::set_bit(bytes, bucket as usize); + } + }; + for part in parts { + match part { + KeyRangePart::Bitmap(part) => { + // Every key below `part.offset` is outside `[min, max]`, so + // this wrapping distance is the part's true position. + let base = part.offset.wrapping_sub(offset); + let width = (1u64 << part.shift) - 1; + for bucket in part.bits.set_indices() { + let lo = (bucket as u64) << part.shift; + let hi = lo.saturating_add(width).min(part.span); + set_range(base + lo, base + hi); + } + } + KeyRangePart::Range { min, max } => { + let lo = ArrayMap::key_to_u64(min)?.wrapping_sub(offset); + let hi = ArrayMap::key_to_u64(max)?.wrapping_sub(offset); + set_range(lo, hi); + } + } + } + + let bits = BooleanBuffer::new(words.into(), 0, buckets as usize); + 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. @@ -214,6 +280,49 @@ mod tests { assert!(kept <= 220, "expected ~200 kept, got {kept}"); } + /// Partitions spanning negative and positive keys, one without a bitmap: + /// the union must keep every container holding a key of any partition, and + /// still prune the gaps no partition covers. + #[test] + fn union_never_drops_a_partition_key() { + let clustered: Vec = (-1_000_000..-999_000).chain(0..1000).collect(); + let sparse: Vec = (0..400).map(|i| 500_000 + 1000 * i).collect(); + let dense = (-200_000i64, -190_000i64); + let a = bitmap(&clustered).expect("not saturated"); + let b = bitmap(&sparse).expect("not saturated"); + let (min, max) = ( + ScalarValue::Int64(Some(-1_000_000)), + ScalarValue::Int64(Some(899_000)), + ); + let (dense_min, dense_max) = ( + ScalarValue::Int64(Some(dense.0)), + ScalarValue::Int64(Some(dense.1)), + ); + let union = KeyRangeBitmap::union( + &min, + &max, + &[ + KeyRangePart::Bitmap(&a), + KeyRangePart::Bitmap(&b), + KeyRangePart::Range { + min: &dense_min, + max: &dense_max, + }, + ], + ) + .expect("not saturated"); + + for key in clustered + .iter() + .chain(&sparse) + .chain(&[dense.0, -195_000, dense.1]) + { + assert!(may(&union, *key, *key), "dropped key {key}"); + } + assert!(!may(&union, -900_000, -800_000)); + assert!(!may(&union, 100_000, 400_000)); + } + #[test] fn contiguous_keys_no_bitmap() { assert!(bitmap(&(0..5000i64).collect::>()).is_none()); From 810516fa3d8680abeb6495a094997022d189e7d6 Mon Sep 17 00:00:00 2001 From: Devan Benz Date: Mon, 19 Oct 2026 18:33:14 -0500 Subject: [PATCH 4/7] Gate internal_err import on proto --- .../src/joins/hash_join/partitioned_hash_eval.rs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) 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 7d184b21a0944..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 @@ -27,6 +27,8 @@ use arrow::{ use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::{create_hashes, with_hashes}; +#[cfg(feature = "proto")] +use datafusion_common::internal_err; use datafusion_expr::ColumnarValue; use datafusion_expr_common::dyn_eq::DynHash; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; @@ -251,9 +253,7 @@ impl HashExpr { let Some(protobuf::physical_expr_node::ExprType::HashExpr(hash_expr)) = &node.expr_type else { - return datafusion_common::internal_err!( - "PhysicalExprNode is not a HashExpr" - ); + return internal_err!("PhysicalExprNode is not a HashExpr"); }; // Destructure exhaustively (no `..`) so that a newly added proto field // is a compile error here instead of being silently ignored. @@ -655,9 +655,7 @@ mod tests { Some(protobuf::physical_expr_node::ExprType::Column(_)) => { Column::try_from_proto(node, &ctx) } - _ => datafusion_common::internal_err!( - "test decoder cannot decode {node:?}" - ), + _ => internal_err!("test decoder cannot decode {node:?}"), } } } From 25b811baec08c6393361b39581c61a084527c55f Mon Sep 17 00:00:00 2001 From: Devan Benz Date: Mon, 19 Oct 2026 19:15:24 -0500 Subject: [PATCH 5/7] Merge InList partitions --- .../physical-plan/src/joins/hash_join/exec.rs | 16 +- .../src/joins/hash_join/shared_bounds.rs | 307 +++++++++++++----- .../physical-plan/src/joins/join_hash_map.rs | 17 +- .../src/joins/key_range_bitmap.rs | 28 +- 4 files changed, 267 insertions(+), 101 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 806cff9773a55..81d167473848a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1864,6 +1864,7 @@ impl ExecutionPlan for HashJoinExec { repartition_random_state, self.null_equality, null_aware, + &context.session_config().options().optimizer, )) }))) }) @@ -10064,19 +10065,8 @@ mod tests { dynamic_filter.wait_complete().await; let filter = dynamic_filter.current()?.to_string(); - // The aggregated build-side NULL must widen the routed filter once, at the - // top: `c2 IS NULL OR CASE hash(c2) % 4 WHEN ... END`. - // - // Under `force_hash_collisions` every key hashes to the same value, so the - // repartition routes all build rows to a single partition and - // `build_partitioned_filter` collapses to that partition's filter with no - // routing `CASE`. The NULL widening this test guards is asserted either way. - if cfg!(not(feature = "force_hash_collisions")) { - assert!( - filter.contains("CASE"), - "expected a routed partitioned filter, got: {filter}" - ); - } + // The aggregated build-side NULL must widen the whole merged filter once, + // at the top: `c2 IS NULL OR `. assert!( filter.starts_with("c2@0 IS NULL OR"), "expected the IS NULL disjunct to wrap the whole routed filter, got: {filter}" 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 1581062cbad67..9049243cd8abc 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -34,12 +34,14 @@ use crate::joins::hash_join::partitioned_hash_eval::{ use crate::joins::join_hash_map::union_hash_maps; use crate::joins::key_range_bitmap::{KeyRangeBitmap, KeyRangePart}; use crate::repartition::RangeExpr; -use arrow::array::ArrayRef; +use arrow::array::{ArrayRef, AsArray}; +use arrow::compute::concat; use arrow::datatypes::{DataType, Field, Schema}; -use datafusion_common::config::ConfigOptions; +use datafusion_common::config::{ConfigOptions, OptimizerOptions}; +use datafusion_common::hash_utils::create_hashes; use datafusion_common::{ DataFusionError, NullEquality, Result, ScalarValue, SharedResult, - assert_or_internal_err, + assert_or_internal_err, internal_err, }; use datafusion_expr::Operator; use datafusion_functions::core::r#struct as struct_func; @@ -294,6 +296,12 @@ pub(crate) struct SharedBuildAccumulator { /// Null-aware anti join (`NOT IN`). A probe-side NULL must reach the join so its /// three-valued logic can collapse the result, so the pushed filter keeps NULL rows. null_aware: Option, + /// Largest merged IN list, in bytes, that partitioned filters keep as an IN + /// list: `hash_join_inlist_pushdown_max_size`, as for a single build. + in_list_max_size: usize, + /// Largest merged IN list, in values, that partitioned filters keep as an IN + /// list: `hash_join_inlist_pushdown_max_distinct_values`, as for a single build. + in_list_max_values: usize, } /// Strategy for filter pushdown (decided at collection time) @@ -406,6 +414,7 @@ impl SharedBuildAccumulator { repartition_random_state: SeededRandomState, null_equality: NullEquality, null_aware: Option, + optimizer: &OptimizerOptions, ) -> Self { // Troubleshooting: If partition counts are incorrect, verify this logic matches // the actual execution pattern in collect_build_side() @@ -463,6 +472,8 @@ impl SharedBuildAccumulator { probe_range_partitioning, null_equality, null_aware, + in_list_max_size: optimizer.hash_join_inlist_pushdown_max_size, + in_list_max_values: optimizer.hash_join_inlist_pushdown_max_distinct_values, } } @@ -688,10 +699,9 @@ impl SharedBuildAccumulator { let mut empty_partition_ids = Vec::new(); let mut has_canceled_unknown = false; let mut keys_have_null = false; - // Every non-empty partition's hash table, pruning bitmap, and bounds, - // merged into one membership lookup below; `None` once a partition's - // build side cannot be merged. - let mut hash_maps = Some(Vec::with_capacity(partitions.len())); + // Every non-empty partition's membership strategy and bounds, merged into + // one membership filter below; `None` once a partition cannot be merged. + let mut members = Some(Vec::with_capacity(partitions.len())); let mut merged_bounds: Option = None; for (partition_id, partition) in partitions.into_iter().enumerate() { @@ -710,19 +720,17 @@ impl SharedBuildAccumulator { }) => { real_partition_ids.push(partition_id); keys_have_null |= partition_keys_have_null; - hash_maps = match (&pushdown, hash_maps) { - (PushdownStrategy::Map(map, bitmap), Some(mut hash_maps)) - if matches!(map.as_ref(), Map::HashMap(_)) => - { - hash_maps.push(( - Arc::clone(map), - bitmap.clone(), - bounds.clone(), - )); - Some(hash_maps) + let mergeable = match &pushdown { + PushdownStrategy::Map(map, _) => { + matches!(map.as_ref(), Map::HashMap(_)) } - _ => None, + PushdownStrategy::InList(_) => true, + PushdownStrategy::Empty => false, }; + members = members.filter(|_| mergeable).map(|mut members| { + members.push((pushdown.clone(), bounds.clone())); + members + }); merged_bounds = match (merged_bounds, real_partition_ids.len()) { (None, 1) => Some(bounds.clone()), (Some(merged), _) => merged.union(&bounds), @@ -744,8 +752,8 @@ impl SharedBuildAccumulator { PartitionStatus::CanceledUnknown => { has_canceled_unknown = true; partition_filters.push(lit(true)); - // Rows routed here must pass, which a merged lookup cannot express. - hash_maps = None; + // Rows routed here must pass, which a merged filter cannot express. + members = None; // A canceled partition's build content is unknown, so it // may hold a NULL key. keys_have_null = true; @@ -775,56 +783,8 @@ impl SharedBuildAccumulator { // Only one build partition contains rows, so its filter covers every // possible probe match without routing. Arc::clone(&partition_filters[real_partition_ids[0]]) - } else if let Some(hash_maps) = hash_maps { - // Every non-empty partition is a hash table. A key is only ever built - // into the partition it is routed to, so membership in the union of - // all tables equals membership in the routed table, and needs neither - // a routing hash nor a CASE over per-partition filters. - let tables = hash_maps - .iter() - .filter_map(|(map, _, _)| match map.as_ref() { - Map::HashMap(table) => Some(table.as_ref()), - Map::ArrayMap(_) => None, - }) - .collect::>(); - let union = Arc::new(Map::HashMap(union_hash_maps(&tables))); - // Container pruning reads one bitmap, so merge the partitions' bitmaps. - // A partition without one may hold a key anywhere in its bounds. Not - // worth building when no partition has a bitmap to contribute gaps. - let pruning_bitmap = merged_bounds - .as_ref() - .and_then(|merged| merged.get_column_bounds(0)) - .filter(|_| { - self.on_right.len() == 1 - && hash_maps.iter().any(|(_, bitmap, _)| bitmap.is_some()) - }) - .and_then(|merged| { - let parts = hash_maps - .iter() - .map(|(_, bitmap, bounds)| match bitmap { - Some(bitmap) => Some(KeyRangePart::Bitmap(bitmap)), - None => bounds.get_column_bounds(0).map(|bounds| { - KeyRangePart::Range { - min: &bounds.min, - max: &bounds.max, - } - }), - }) - .collect::>>()?; - KeyRangeBitmap::union(&merged.min, &merged.max, &parts) - }) - .map(Arc::new); - let membership = Arc::new(HashTableLookupExpr::new( - self.on_right.clone(), - HASH_JOIN_SEED, - union, - String::from("hash_lookup"), - pruning_bitmap, - )) as Arc; - let bounds = merged_bounds - .and_then(|bounds| create_bounds_predicate(&self.on_right, &bounds)); - combine_membership_and_bounds(Some(membership), bounds) - .expect("membership predicate is always present") + } else if let Some(members) = members { + self.merged_membership_filter(&members, merged_bounds)? } else { // Builds the shared sparse `CASE` for partition filter routing. // Without cancellation, omitted branches are known empty and safely fall @@ -899,6 +859,126 @@ impl SharedBuildAccumulator { .update(self.preserve_probe_nulls(filter_expr, keys_have_null)?) } + /// One membership filter over every partition's build keys, replacing the + /// `CASE` over per-partition filters. A key is only ever built into the + /// partition it is routed to, so membership in the union of all partitions' + /// keys equals membership in the routed partition's keys, and needs neither + /// a routing hash nor per-partition branches. + /// + /// `members` holds each non-empty partition's strategy, which must be an IN + /// list or a hash table, and bounds. + fn merged_membership_filter( + &self, + members: &[(PushdownStrategy, PartitionBounds)], + merged_bounds: Option, + ) -> Result> { + let mut tables = Vec::with_capacity(members.len()); + let mut lists = Vec::with_capacity(members.len()); + for (pushdown, _) in members { + match pushdown { + PushdownStrategy::Map(map, _) => match map.as_ref() { + Map::HashMap(table) => tables.push(table.as_ref()), + Map::ArrayMap(_) => { + return internal_err!("array maps cannot be merged"); + } + }, + PushdownStrategy::InList(values) => lists.push(values), + PushdownStrategy::Empty => {} + } + } + + // Every partition is an IN list, and together they still fit the limits a + // single build uses to pick an IN list: one combined list keeps the + // statistics and bloom filter pruning an IN list allows. Past those + // limits, a hash lookup is cheaper per row and carries a pruning bitmap. + let combined_in_list = tables.is_empty() && { + let values = lists.iter().map(|list| list.len()).sum::(); + let bytes = lists + .iter() + .map(|list| list.get_array_memory_size()) + .sum::(); + values <= self.in_list_max_values && bytes <= self.in_list_max_size + }; + let membership = if combined_in_list { + let values = concat(&lists.iter().map(|l| l.as_ref()).collect::>())?; + create_membership_predicate( + &self.on_right, + PushdownStrategy::InList(values), + &HASH_JOIN_SEED, + self.probe_schema.as_ref(), + )? + .expect("an IN list always has a membership predicate") + } else { + // IN list keys join the hash tables' keys, hashed as the probe side + // hashes its join keys: one column per join key. + let mut hashes = vec![0; lists.iter().map(|l| l.len()).sum()]; + let mut start = 0; + for values in &lists { + let columns = match values.as_struct_opt() { + Some(columns) if self.on_right.len() > 1 => columns.columns(), + _ => std::slice::from_ref(*values), + }; + let end = start + values.len(); + create_hashes( + columns, + HASH_JOIN_SEED.random_state(), + &mut hashes[start..end], + )?; + start = end; + } + let union = Arc::new(Map::HashMap(union_hash_maps(&tables, &hashes))); + Arc::new(HashTableLookupExpr::new( + self.on_right.clone(), + HASH_JOIN_SEED, + union, + String::from("hash_lookup"), + self.merged_pruning_bitmap(members, merged_bounds.as_ref()), + )) + }; + + let bounds = merged_bounds + .and_then(|bounds| create_bounds_predicate(&self.on_right, &bounds)); + Ok(combine_membership_and_bounds(Some(membership), bounds) + .expect("membership predicate is always present")) + } + + /// Container pruning reads one bitmap, so merge the partitions' key ranges: + /// a hash table's bitmap, an IN list's keys, or else a partition's whole + /// bounds. `None` when no partition narrows its bounds, as the bounds + /// predicate already prunes that much. + fn merged_pruning_bitmap( + &self, + members: &[(PushdownStrategy, PartitionBounds)], + merged_bounds: Option<&PartitionBounds>, + ) -> Option> { + if self.on_right.len() != 1 { + return None; + } + let merged = merged_bounds?.get_column_bounds(0)?; + let parts = members + .iter() + .map(|(pushdown, bounds)| match pushdown { + PushdownStrategy::Map(_, Some(bitmap)) => { + Some(KeyRangePart::Bitmap(bitmap)) + } + PushdownStrategy::InList(keys) => Some(KeyRangePart::Keys(keys.as_ref())), + _ => bounds + .get_column_bounds(0) + .map(|bounds| KeyRangePart::Range { + min: &bounds.min, + max: &bounds.max, + }), + }) + .collect::>>()?; + if parts + .iter() + .all(|part| matches!(part, KeyRangePart::Range { .. })) + { + return None; + } + KeyRangeBitmap::union(&merged.min, &merged.max, &parts).map(Arc::new) + } + /// Keeps probe rows with a NULL key when the join semantics need them. /// /// The build-side predicate drops probe rows whose key is NULL. A null-aware join @@ -995,6 +1075,8 @@ pub(super) fn make_partitioned_accumulator_for_test( probe_range_partitioning: None, null_equality: NullEquality::NullEqualsNothing, null_aware: None, + in_list_max_size: 128 * 1024, + in_list_max_values: 150, } } @@ -1061,6 +1143,8 @@ mod tests { probe_range_partitioning: None, null_equality: NullEquality::NullEqualsNothing, null_aware: None, + in_list_max_size: 128 * 1024, + in_list_max_values: 150, } } @@ -1258,17 +1342,80 @@ mod tests { ); } + /// A hash table holding `keys`, hashed as the build side hashes them. + #[cfg(not(feature = "force_hash_collisions"))] + fn hash_map_of(keys: &[i32]) -> PushdownStrategy { + use crate::joins::join_hash_map::JoinHashMapType; + let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec())); + let mut hashes = vec![0; keys.len()]; + create_hashes([&keys], HASH_JOIN_SEED.random_state(), &mut hashes).unwrap(); + let mut map = JoinHashMapU32::with_capacity(keys.len()); + map.update_from_iter(Box::new(hashes.iter().enumerate()), 0); + PushdownStrategy::Map(Arc::new(Map::HashMap(Box::new(map))), None) + } + + /// Partitions that chose an IN list and partitions that chose a hash table + /// merge into one lookup without routing: rows holding any partition's key + /// pass, and rows holding none are filtered. #[test] - fn partitioned_in_list_partition_keeps_case() { - let acc = make_partitioned_expr_accumulator_for_test(2); + #[cfg(not(feature = "force_hash_collisions"))] + fn partitioned_in_list_and_hash_map_partitions_merge() { + let acc = make_partitioned_expr_accumulator_for_test(3); acc.build_filter(FinalizeInput::Partitioned(vec![ - reported(hash_map(), no_bounds()), + reported(hash_map_of(&[7, 9]), no_bounds()), reported(in_list(&[2]), no_bounds()), + reported(PushdownStrategy::Empty, no_bounds()), ])) .unwrap(); - assert!(current_expr(&acc).downcast_ref::().is_some()); + assert!( + current_expr(&acc).downcast_ref::().is_none(), + "expected a merged lookup, got: {}", + current_expr(&acc) + ); + assert_matches(&acc, &[2, 7, 9, 3, 8], &[true, true, true, false, false]); + } + + /// IN list partitions whose combined list passes the IN list limits merge + /// into a hash lookup rather than one oversized IN list, which must still + /// accept every partition's keys. + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn partitioned_in_lists_past_limit_merge_into_hash_lookup() { + let mut acc = make_partitioned_expr_accumulator_for_test(3); + acc.in_list_max_values = 2; + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(in_list(&[1]), no_bounds()), + reported(in_list(&[2]), no_bounds()), + reported(in_list(&[3]), no_bounds()), + ])) + .unwrap(); + + let expr = current_expr(&acc); + assert!( + expr.downcast_ref::().is_none() + && expr.downcast_ref::().is_none(), + "expected a hash lookup, got: {expr}" + ); + assert_matches(&acc, &[1, 2, 3, 4], &[true, true, true, false]); + } + + /// Evaluates the current filter over probe keys `keys`. + #[cfg(not(feature = "force_hash_collisions"))] + fn assert_matches(acc: &SharedBuildAccumulator, keys: &[i32], expected: &[bool]) { + let probe = RecordBatch::try_new( + test_probe_schema(), + vec![Arc::new(Int32Array::from(keys.to_vec()))], + ) + .unwrap(); + let matched = current_expr(acc) + .evaluate(&probe) + .unwrap() + .into_array(keys.len()) + .unwrap(); + assert_eq!(matched.as_boolean(), &BooleanArray::from(expected.to_vec())); } #[test] @@ -1591,6 +1738,8 @@ mod tests { correlated: false, value_keys: 1, }), + in_list_max_size: 128 * 1024, + in_list_max_values: 150, } } @@ -1638,9 +1787,9 @@ mod tests { acc } - /// The IS NULL disjunct must wrap the routed `CASE` as a whole. A probe NULL - /// takes exactly one route, so widening only the branch of the partition that - /// reported the NULL would still drop it on every other route. + /// The IS NULL disjunct must wrap the merged filter as a whole. A probe NULL + /// belongs to exactly one partition, so widening only the filter of the + /// partition that reported the NULL would still drop it everywhere else. #[test] fn partitioned_null_keys_in_one_partition_widen_whole_routed_filter() { let acc = null_equal_partitioned_accumulator(2); @@ -1658,7 +1807,7 @@ mod tests { widened.left().downcast_ref::().is_some(), "expected the IS NULL disjunct first, got: {expr}" ); - case_expr(widened.right()); + assert_in_list_column_values(widened.right(), "probe_key", 0, &[1, 2]); assert_eq!( format!("{expr}").matches("IS NULL").count(), 1, @@ -1806,6 +1955,8 @@ mod tests { correlated: true, value_keys, }), + in_list_max_size: 128 * 1024, + in_list_max_values: 150, }; let bounds = PartitionBounds::new( diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index 1a022c4283d56..a20cfdf4f5e05 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -410,15 +410,22 @@ where } } -/// Builds one map holding the union of `maps`' hash values, so a single -/// [`JoinHashMapType::contain_hashes`] call checks membership in all of them. +/// Builds one map holding the union of `maps`' hash values and `hashes`, so a +/// single [`JoinHashMapType::contain_hashes`] call checks membership in all of +/// them. /// /// Its row indices number the merged hash values rather than rows of any build /// batch, so it is only meaningful for membership checks, never to produce /// join matches. -pub(crate) fn union_hash_maps(maps: &[&dyn JoinHashMapType]) -> Box { - let len = maps.iter().map(|map| map.len()).sum(); - let hashes = maps.iter().flat_map(|map| map.hash_values()); +pub(crate) fn union_hash_maps( + maps: &[&dyn JoinHashMapType], + hashes: &[u64], +) -> Box { + let len = maps.iter().map(|map| map.len()).sum::() + hashes.len(); + let hashes = maps + .iter() + .flat_map(|map| map.hash_values()) + .chain(hashes.iter().copied()); if len > u32::MAX as usize { let mut union = JoinHashMapU64::with_capacity(len); for (row, hash) in hashes.enumerate() { diff --git a/datafusion/physical-plan/src/joins/key_range_bitmap.rs b/datafusion/physical-plan/src/joins/key_range_bitmap.rs index bd87db415b558..6ef4834294fd2 100644 --- a/datafusion/physical-plan/src/joins/key_range_bitmap.rs +++ b/datafusion/physical-plan/src/joins/key_range_bitmap.rs @@ -32,6 +32,8 @@ const MAX_BUCKETS: u64 = 1 << 20; pub(crate) enum KeyRangePart<'a> { /// The partition's own bitmap. Bitmap(&'a KeyRangeBitmap), + /// The partition's build keys, which must lie within its `[min, max]`. + Keys(&'a dyn Array), /// The partition has no bitmap, so any key in `[min, max]` may be present. Range { min: &'a ScalarValue, @@ -131,7 +133,7 @@ impl KeyRangeBitmap { let mut words = MutableBuffer::new_null(buckets as usize); let bytes = words.as_slice_mut(); // Sets the buckets covering keys `lo..=hi`, given as offsets from `offset`. - let mut set_range = |lo: u64, hi: u64| { + let set_range = |bytes: &mut [u8], lo: u64, hi: u64| { for bucket in (lo >> shift)..=(hi.min(span) >> shift) { bit_util::set_bit(bytes, bucket as usize); } @@ -146,13 +148,19 @@ impl KeyRangeBitmap { for bucket in part.bits.set_indices() { let lo = (bucket as u64) << part.shift; let hi = lo.saturating_add(width).min(part.span); - set_range(base + lo, base + hi); + set_range(bytes, base + lo, base + hi); } } + KeyRangePart::Keys(keys) => { + if !ArrayMap::is_supported_type(keys.data_type()) { + return None; + } + set_bucket_bits(*keys, offset, shift, buckets, bytes); + } KeyRangePart::Range { min, max } => { let lo = ArrayMap::key_to_u64(min)?.wrapping_sub(offset); let hi = ArrayMap::key_to_u64(max)?.wrapping_sub(offset); - set_range(lo, hi); + set_range(bytes, lo, hi); } } } @@ -223,9 +231,20 @@ impl KeyRangeBitmap { /// 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); + set_bucket_bits(keys, offset, shift, buckets, words.as_slice_mut()); + BooleanBuffer::new(words.into(), 0, buckets as usize) +} + +/// Sets the bit of each bucket holding a non-null key of `keys` in `bytes`. +fn set_bucket_bits( + keys: &dyn Array, + offset: u64, + shift: u32, + buckets: u64, + bytes: &mut [u8], +) { 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); @@ -238,7 +257,6 @@ fn bucket_bits(keys: &dyn Array, offset: u64, shift: u32, buckets: u64) -> Boole } _ => unreachable!("guarded by ArrayMap::is_supported_type"), ); - BooleanBuffer::new(words.into(), 0, buckets as usize) } #[cfg(test)] From e52d70849aca0dabb39312547583ef5057257ae2 Mon Sep 17 00:00:00 2001 From: Devan Benz Date: Tue, 20 Oct 2026 13:46:47 -0500 Subject: [PATCH 6/7] chore(tests): fix tests --- .../tests/physical_optimizer/filter_pushdown.rs | 4 ++-- .../src/joins/hash_join/shared_bounds.rs | 15 ++++++++++++--- .../test_files/range_partitioning.slt | 9 ++++----- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index ddd4bb4d271aa..67f1dcdc10afa 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -1165,7 +1165,7 @@ async fn test_hashjoin_dynamic_filter_pushdown_partitioned() { - RepartitionExec: partitioning=Hash([a@0, b@1], 12), input_partitions=1 - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true - RepartitionExec: partitioning=Hash([a@0, b@1], 12), input_partitions=1 - - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ CASE hash_repartition % 12 WHEN 5 THEN a@0 >= ab AND a@0 <= ab AND b@1 >= bb AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}]) WHEN 8 THEN a@0 >= aa AND a@0 <= aa AND b@1 >= ba AND b@1 <= ba AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}]) ELSE false END ] + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}, {c0:aa,c1:ba}]) ] " ); @@ -1381,7 +1381,7 @@ async fn test_hashjoin_dynamic_filter_pushdown_range_partitioned() { - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 - - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN a@0 >= aa AND a@0 <= aa AND b@1 >= ba AND b@1 <= ba AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}]) WHEN 1 THEN a@0 >= ab AND a@0 <= ab AND b@1 >= bb AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}]) ELSE false END ] + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ] " ); 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 9049243cd8abc..dde71a9004771 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -90,7 +90,8 @@ impl PartitionBounds { } /// Bounds covering both `self` and `other`, or `None` if they do not - /// describe the same columns. + /// describe the same columns. A NULL bound means the partition holds no + /// non-null key in that column, so it never widens the other side's bound. fn union(self, other: &PartitionBounds) -> Option { if self.column_bounds.len() != other.column_bounds.len() { return None; @@ -100,8 +101,16 @@ impl PartitionBounds { .into_iter() .zip(&other.column_bounds) .map(|(a, b)| { - let min = if b.min < a.min { b.min.clone() } else { a.min }; - let max = if b.max > a.max { b.max.clone() } else { a.max }; + let min = if !b.min.is_null() && (a.min.is_null() || b.min < a.min) { + b.min.clone() + } else { + a.min + }; + let max = if !b.max.is_null() && (a.max.is_null() || b.max > a.max) { + b.max.clone() + } else { + a.max + }; ColumnBounds::new(min, max) }) .collect(); diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 56184845a6fe3..9f773197a97f3 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -1899,10 +1899,9 @@ SELECT range_key, SUM(value) FROM ( ########## # TEST 48: Hash Join Dynamic Filter Pushdown on Compatible Range Inputs -# The Parquet-backed probe accepts the partition-routed dynamic filter. Matching -# Range split points keep build filter i aligned with probe partition i. The -# build has rows only in partitions 0 and 2, so the sparse runtime CASE includes -# only those branches and lets empty partitions fall through to false. +# The Parquet-backed probe accepts the dynamic filter. The build has rows only in +# partitions 0 and 2, and both chose an IN list, so they merge into one IN list +# with combined bounds instead of a per-partition CASE. ########## statement ok @@ -1943,7 +1942,7 @@ JOIN range_partitioned p ON b.range_key = p.range_key; Plan with Metrics 01)HashJoinExec: mode=Partitionedmetrics=[output_rows=2,] 02)--DataSourceExec: file_type=parquet, predicate=range_key@0 = 5 OR range_key@0 = 20metrics=[output_rows=2,] -03)--DataSourceExec: file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN range_key@0 >= 5 AND range_key@0 <= 5 AND range_key@0 IN (SET) ([5]) WHEN 2 THEN range_key@0 >= 20 AND range_key@0 <= 20 AND range_key@0 IN (SET) ([20]) ELSE false END ]metrics=[output_rows=2,] +03)--DataSourceExec: file_type=parquet, predicate=DynamicFilter [ range_key@0 >= 5 AND range_key@0 <= 20 AND range_key@0 IN (SET) ([5, 20]) ]metrics=[output_rows=2,] query III SELECT b.range_key, b.value, p.value From 160548dd3e18850a836578ad11f5529696f5cc7b Mon Sep 17 00:00:00 2001 From: Devan Benz Date: Tue, 20 Oct 2026 14:54:11 -0500 Subject: [PATCH 7/7] feat(perf): route merged hash lookups instead of building a union table Building one hash table over every build partition's keys ran on a single thread at the dynamic filter barrier, so every probe stream waited on it (about 0.7s for a 15M key build at TPC-H SF10). HashTableLookupExpr now has a routed mode: each probe row is looked up in the table of the build partition its key was repartitioned to (hash or range routing), sharing the partitions' tables as they are. It stays a HashTableLookupExpr, so pruning keeps using its merged pruning bitmap. Removes union_hash_maps and JoinHashMapType::hash_values. --- .../joins/hash_join/partitioned_hash_eval.rs | 213 ++++++++++++++++-- .../src/joins/hash_join/shared_bounds.rs | 203 +++++++++++------ .../physical-plan/src/joins/join_hash_map.rs | 100 ++------ .../src/joins/stream_join_utils.rs | 4 - .../physical-plan/src/repartition/mod.rs | 14 ++ 5 files changed, 366 insertions(+), 168 deletions(-) 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 921a38970562e..293c6e932ce2d 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 @@ -20,21 +20,20 @@ use std::{fmt::Display, hash::Hash, sync::Arc}; use arrow::{ - array::{ArrayRef, UInt64Array}, + array::{ArrayRef, BooleanArray, BooleanBufferBuilder, UInt64Array}, datatypes::{DataType, Schema}, record_batch::RecordBatch, }; -use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::{create_hashes, with_hashes}; -#[cfg(feature = "proto")] -use datafusion_common::internal_err; +use datafusion_common::{Result, internal_err}; use datafusion_expr::ColumnarValue; 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; +use crate::repartition::RangeExpr; /// RandomState wrapper that preserves the seed used to create it. /// @@ -281,13 +280,71 @@ pub struct HashTableLookupExpr { on_columns: Vec, /// Random state for hashing (with seeds preserved for serialization) random_state: SeededRandomState, - /// Map to check against (hash table or array map) - map: Arc, + /// Map(s) to check against + lookup: Lookup, /// Description for display description: String, /// Shared pruning-only view of the build side's key range pruning_bitmap: Option>, } + +/// The build side a [`HashTableLookupExpr`] checks probe keys against. +#[derive(Clone)] +enum Lookup { + /// One map holding every build key. + Single(Arc), + /// One hash table per build partition of a `Partitioned` join, `None` where + /// the partition is empty. A key is only ever built into the partition its + /// key routes to, so each probe row is looked up in that partition's table + /// alone: no `CASE` over per-partition filters, and no merged table. + Routed { + router: PartitionRouter, + maps: Arc<[Option>]>, + }, +} + +/// Assigns probe rows to the build partition their key was repartitioned to. +#[derive(Clone)] +pub(crate) enum PartitionRouter { + /// `hash(keys) % partitions`, as hash repartitioning routes rows. + Hash(SeededRandomState), + /// The probe side's range partitioning. + Range(Arc), +} + +impl PartitionRouter { + /// Writes the partition of each row of `keys` to `partition_ids`, which + /// holds one slot per row. + fn route( + &self, + keys: &[ArrayRef], + partitions: usize, + partition_ids: &mut Vec, + ) -> Result<()> { + match self { + PartitionRouter::Hash(random_state) => { + partition_ids.resize(keys[0].len(), 0); + create_hashes(keys, random_state.random_state(), partition_ids)?; + let partitions = partitions as u64; + for id in partition_ids.iter_mut() { + *id %= partitions; + } + Ok(()) + } + PartitionRouter::Range(range) => { + range.route_partition_ids(keys, partition_ids) + } + } + } + + fn seed(&self) -> Option { + match self { + PartitionRouter::Hash(random_state) => Some(random_state.seed()), + PartitionRouter::Range(_) => None, + } + } +} + impl HashTableLookupExpr { /// Create a new HashTableLookupExpr /// @@ -312,16 +369,101 @@ impl HashTableLookupExpr { Self { on_columns, random_state, - map, + lookup: Lookup::Single(map), description, pruning_bitmap, } } + /// Looks each probe row up in the hash table of the build partition + /// `router` assigns it to. `maps` holds one entry per partition, each a + /// [`Map::HashMap`] or `None` for an empty partition, and `pruning_bitmap` + /// must cover every partition's keys. + pub(crate) fn new_routed( + on_columns: Vec, + random_state: SeededRandomState, + router: PartitionRouter, + maps: Vec>>, + description: String, + pruning_bitmap: Option>, + ) -> Result { + if maps + .iter() + .flatten() + .any(|map| !matches!(map.as_ref(), Map::HashMap(_))) + { + return internal_err!("routed lookups only support hash table partitions"); + } + Ok(Self { + on_columns, + random_state, + lookup: Lookup::Routed { + router, + maps: maps.into(), + }, + 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() } + + /// Membership of each row of `keys` in the hash table of its routed partition. + fn evaluate_routed( + &self, + keys: &[ArrayRef], + router: &PartitionRouter, + maps: &[Option>], + ) -> Result { + let num_rows = keys[0].len(); + let mut partition_ids = Vec::with_capacity(num_rows); + router.route(keys, maps.len(), &mut partition_ids)?; + + with_hashes(keys, self.random_state.random_state(), |hashes| { + // Counting sort of the rows by partition, so each table is probed + // once per batch over a contiguous run of its rows' hashes. + let mut starts = vec![0usize; maps.len() + 1]; + for &id in &partition_ids { + starts[id as usize + 1] += 1; + } + for i in 0..maps.len() { + starts[i + 1] += starts[i]; + } + let mut next = starts.clone(); + let mut rows = vec![0u32; num_rows]; + let mut sorted_hashes = vec![0u64; num_rows]; + for (row, (&id, &hash)) in partition_ids.iter().zip(hashes).enumerate() { + let slot = &mut next[id as usize]; + rows[*slot] = row as u32; + sorted_hashes[*slot] = hash; + *slot += 1; + } + + let mut matched = BooleanBufferBuilder::new(num_rows); + matched.append_n(num_rows, false); + for (partition, map) in maps.iter().enumerate() { + // An empty partition holds no key, so its rows stay unmatched. + let Some(map) = map else { continue }; + let Map::HashMap(map) = map.as_ref() else { + return internal_err!( + "routed lookups only support hash table partitions" + ); + }; + let run = starts[partition]..starts[partition + 1]; + if run.is_empty() { + continue; + } + let hits = map.contain_hashes(&sorted_hashes[run.clone()]); + for hit in hits.values().set_indices() { + matched.set_bit(rows[run.start + hit] as usize, true); + } + } + Ok(BooleanArray::new(matched.finish(), None)) + }) + } } impl std::fmt::Debug for HashTableLookupExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -332,7 +474,21 @@ impl std::fmt::Debug for HashTableLookupExpr { .collect::>() .join(", "); let seed = self.random_state.seed(); - write!(f, "{}({cols}, [{seed}])", self.description) + match &self.lookup { + Lookup::Single(_) => write!(f, "{}({cols}, [{seed}])", self.description), + Lookup::Routed { router, maps } => { + let router = match router.seed() { + Some(routing_seed) => format!("hash [{routing_seed}]"), + None => "range".to_string(), + }; + write!( + f, + "{}({cols}, [{seed}], routed by {router} over {} partitions)", + self.description, + maps.len() + ) + } + } } } @@ -342,7 +498,7 @@ impl Hash for HashTableLookupExpr { self.description.hash(state); self.random_state.seed().hash(state); // Hash immutable map identity, consistently with pointer equality below. - Arc::as_ptr(&self.map).hash(state); + self.lookup.identity().hash(state); } } @@ -353,12 +509,23 @@ impl PartialEq for HashTableLookupExpr { self.on_columns == other.on_columns && self.description == other.description && self.random_state.seed() == other.random_state.seed() - && Arc::ptr_eq(&self.map, &other.map) + && self.lookup.identity() == other.lookup.identity() } } impl Eq for HashTableLookupExpr {} +impl Lookup { + /// Identity of the immutable build-side maps. A routed lookup's router is + /// created together with its maps, so the maps identify it as well. + fn identity(&self) -> *const () { + match self { + Lookup::Single(map) => Arc::as_ptr(map).cast(), + Lookup::Routed { maps, .. } => Arc::as_ptr(maps).cast(), + } + } +} + impl Display for HashTableLookupExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.description) @@ -377,7 +544,7 @@ impl PhysicalExpr for HashTableLookupExpr { Ok(Arc::new(HashTableLookupExpr { on_columns: children, random_state: self.random_state.clone(), - map: Arc::clone(&self.map), + lookup: self.lookup.clone(), description: self.description.clone(), pruning_bitmap: self.pruning_bitmap.clone(), })) @@ -395,15 +562,21 @@ impl PhysicalExpr for HashTableLookupExpr { // Evaluate columns let join_keys = evaluate_columns(&self.on_columns, batch)?; - match self.map.as_ref() { - Map::HashMap(map) => { - with_hashes(&join_keys, self.random_state.random_state(), |hashes| { - let array = map.contain_hashes(hashes); + match &self.lookup { + Lookup::Single(map) => match map.as_ref() { + Map::HashMap(map) => { + with_hashes(&join_keys, self.random_state.random_state(), |hashes| { + let array = map.contain_hashes(hashes); + Ok(ColumnarValue::Array(Arc::new(array))) + }) + } + Map::ArrayMap(map) => { + let array = map.contain_keys(&join_keys)?; Ok(ColumnarValue::Array(Arc::new(array))) - }) - } - Map::ArrayMap(map) => { - let array = map.contain_keys(&join_keys)?; + } + }, + Lookup::Routed { router, maps } => { + let array = self.evaluate_routed(&join_keys, router, maps)?; Ok(ColumnarValue::Array(Arc::new(array))) } } @@ -424,7 +597,7 @@ impl PhysicalExpr for HashTableLookupExpr { // with lit(true), see the comment below on_columns: _, random_state: _, - map: _, + lookup: _, description: _, pruning_bitmap: _, } = self; 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 dde71a9004771..0d63e67480c3d 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -29,9 +29,9 @@ use crate::joins::PartitionMode; use crate::joins::hash_join::exec::{HASH_JOIN_SEED, NullAwareMode}; use crate::joins::hash_join::inlist_builder::build_struct_fields; use crate::joins::hash_join::partitioned_hash_eval::{ - HashExpr, HashTableLookupExpr, SeededRandomState, + HashExpr, HashTableLookupExpr, PartitionRouter, SeededRandomState, }; -use crate::joins::join_hash_map::union_hash_maps; +use crate::joins::join_hash_map::{JoinHashMapType, JoinHashMapU32}; use crate::joins::key_range_bitmap::{KeyRangeBitmap, KeyRangePart}; use crate::repartition::RangeExpr; use arrow::array::{ArrayRef, AsArray}; @@ -737,7 +737,7 @@ impl SharedBuildAccumulator { PushdownStrategy::Empty => false, }; members = members.filter(|_| mergeable).map(|mut members| { - members.push((pushdown.clone(), bounds.clone())); + members.push((partition_id, pushdown.clone(), bounds.clone())); members }); merged_bounds = match (merged_bounds, real_partition_ids.len()) { @@ -793,7 +793,11 @@ impl SharedBuildAccumulator { // possible probe match without routing. Arc::clone(&partition_filters[real_partition_ids[0]]) } else if let Some(members) = members { - self.merged_membership_filter(&members, merged_bounds)? + self.merged_membership_filter( + &members, + merged_bounds, + partition_filters.len(), + )? } else { // Builds the shared sparse `CASE` for partition filter routing. // Without cancellation, omitted branches are known empty and safely fall @@ -869,38 +873,29 @@ impl SharedBuildAccumulator { } /// One membership filter over every partition's build keys, replacing the - /// `CASE` over per-partition filters. A key is only ever built into the - /// partition it is routed to, so membership in the union of all partitions' - /// keys equals membership in the routed partition's keys, and needs neither - /// a routing hash nor per-partition branches. + /// `CASE` over per-partition filters. /// - /// `members` holds each non-empty partition's strategy, which must be an IN - /// list or a hash table, and bounds. + /// `members` holds each non-empty partition's id, strategy (an IN list or a + /// hash table) and bounds, out of `num_partitions` partitions. fn merged_membership_filter( &self, - members: &[(PushdownStrategy, PartitionBounds)], + members: &[(usize, PushdownStrategy, PartitionBounds)], merged_bounds: Option, + num_partitions: usize, ) -> Result> { - let mut tables = Vec::with_capacity(members.len()); - let mut lists = Vec::with_capacity(members.len()); - for (pushdown, _) in members { - match pushdown { - PushdownStrategy::Map(map, _) => match map.as_ref() { - Map::HashMap(table) => tables.push(table.as_ref()), - Map::ArrayMap(_) => { - return internal_err!("array maps cannot be merged"); - } - }, - PushdownStrategy::InList(values) => lists.push(values), - PushdownStrategy::Empty => {} - } - } + let lists = members + .iter() + .filter_map(|(_, pushdown, _)| match pushdown { + PushdownStrategy::InList(values) => Some(values), + _ => None, + }) + .collect::>(); // Every partition is an IN list, and together they still fit the limits a // single build uses to pick an IN list: one combined list keeps the // statistics and bloom filter pruning an IN list allows. Past those // limits, a hash lookup is cheaper per row and carries a pruning bitmap. - let combined_in_list = tables.is_empty() && { + let combined_in_list = lists.len() == members.len() && { let values = lists.iter().map(|list| list.len()).sum::(); let bytes = lists .iter() @@ -918,31 +913,27 @@ impl SharedBuildAccumulator { )? .expect("an IN list always has a membership predicate") } else { - // IN list keys join the hash tables' keys, hashed as the probe side - // hashes its join keys: one column per join key. - let mut hashes = vec![0; lists.iter().map(|l| l.len()).sum()]; - let mut start = 0; - for values in &lists { - let columns = match values.as_struct_opt() { - Some(columns) if self.on_right.len() > 1 => columns.columns(), - _ => std::slice::from_ref(*values), - }; - let end = start + values.len(); - create_hashes( - columns, - HASH_JOIN_SEED.random_state(), - &mut hashes[start..end], - )?; - start = end; + // Each probe row is looked up in the table of the partition its key + // routes to, so the partitions' tables are shared as they are: no + // merged table to build while every probe stream waits on the filter. + let mut maps = vec![None; num_partitions]; + for (partition_id, pushdown, _) in members { + maps[*partition_id] = Some(match pushdown { + PushdownStrategy::Map(map, _) => Arc::clone(map), + PushdownStrategy::InList(values) => self.in_list_hash_map(values)?, + PushdownStrategy::Empty => { + return internal_err!("empty partitions have no members"); + } + }); } - let union = Arc::new(Map::HashMap(union_hash_maps(&tables, &hashes))); - Arc::new(HashTableLookupExpr::new( + Arc::new(HashTableLookupExpr::new_routed( self.on_right.clone(), HASH_JOIN_SEED, - union, + self.partition_router(num_partitions)?, + maps, String::from("hash_lookup"), self.merged_pruning_bitmap(members, merged_bounds.as_ref()), - )) + )?) }; let bounds = merged_bounds @@ -951,13 +942,48 @@ impl SharedBuildAccumulator { .expect("membership predicate is always present")) } + /// How probe rows find their build partition: as the probe side was + /// partitioned. + fn partition_router(&self, num_partitions: usize) -> Result { + Ok(match &self.probe_range_partitioning { + Some(range_partitioning) => { + assert_or_internal_err!( + num_partitions == range_partitioning.partition_count(), + "Dynamic filter partition count {} does not match Range partition count {}", + num_partitions, + range_partitioning.partition_count() + ); + PartitionRouter::Range(Arc::new(RangeExpr::try_new_with_schema( + self.on_right.clone(), + range_partitioning, + &self.probe_schema, + )?)) + } + None => PartitionRouter::Hash(self.repartition_random_state.clone()), + }) + } + + /// A hash table of an IN list partition's keys, hashed as the probe side + /// hashes its join keys: one column per join key. + fn in_list_hash_map(&self, values: &ArrayRef) -> Result> { + let columns = match values.as_struct_opt() { + Some(columns) if self.on_right.len() > 1 => columns.columns(), + _ => std::slice::from_ref(values), + }; + let mut hashes = vec![0; values.len()]; + create_hashes(columns, HASH_JOIN_SEED.random_state(), &mut hashes)?; + let mut map = JoinHashMapU32::with_capacity(hashes.len()); + map.update_from_iter(Box::new(hashes.iter().enumerate()), 0); + Ok(Arc::new(Map::HashMap(Box::new(map)))) + } + /// Container pruning reads one bitmap, so merge the partitions' key ranges: /// a hash table's bitmap, an IN list's keys, or else a partition's whole /// bounds. `None` when no partition narrows its bounds, as the bounds /// predicate already prunes that much. fn merged_pruning_bitmap( &self, - members: &[(PushdownStrategy, PartitionBounds)], + members: &[(usize, PushdownStrategy, PartitionBounds)], merged_bounds: Option<&PartitionBounds>, ) -> Option> { if self.on_right.len() != 1 { @@ -966,7 +992,7 @@ impl SharedBuildAccumulator { let merged = merged_bounds?.get_column_bounds(0)?; let parts = members .iter() - .map(|(pushdown, bounds)| match pushdown { + .map(|(_, pushdown, bounds)| match pushdown { PushdownStrategy::Map(_, Some(bitmap)) => { Some(KeyRangePart::Bitmap(bitmap)) } @@ -1363,19 +1389,62 @@ mod tests { PushdownStrategy::Map(Arc::new(Map::HashMap(Box::new(map))), None) } + /// Splits `keys` across `num_partitions` build partitions as hash + /// repartitioning would, reporting each non-empty partition with the + /// strategy `strategy` picks for it. + #[cfg(not(feature = "force_hash_collisions"))] + fn routed_partitions( + acc: &SharedBuildAccumulator, + num_partitions: usize, + keys: &[i32], + strategy: impl Fn(usize, &[i32]) -> PushdownStrategy, + ) -> Vec { + let array: ArrayRef = Arc::new(Int32Array::from(keys.to_vec())); + let mut hashes = vec![0; keys.len()]; + create_hashes( + [&array], + acc.repartition_random_state.random_state(), + &mut hashes, + ) + .unwrap(); + (0..num_partitions) + .map(|partition| { + let held: Vec = keys + .iter() + .zip(&hashes) + .filter(|(_, hash)| *hash % num_partitions as u64 == partition as u64) + .map(|(key, _)| *key) + .collect(); + if held.is_empty() { + reported(PushdownStrategy::Empty, no_bounds()) + } else { + reported(strategy(partition, &held), no_bounds()) + } + }) + .collect() + } + /// Partitions that chose an IN list and partitions that chose a hash table - /// merge into one lookup without routing: rows holding any partition's key - /// pass, and rows holding none are filtered. + /// merge into one lookup without a `CASE`: a probe row passes exactly when + /// the partition its key routes to holds the key. #[test] #[cfg(not(feature = "force_hash_collisions"))] fn partitioned_in_list_and_hash_map_partitions_merge() { - let acc = make_partitioned_expr_accumulator_for_test(3); - - acc.build_filter(FinalizeInput::Partitioned(vec![ - reported(hash_map_of(&[7, 9]), no_bounds()), - reported(in_list(&[2]), no_bounds()), - reported(PushdownStrategy::Empty, no_bounds()), - ])) + let acc = make_partitioned_expr_accumulator_for_test(4); + let build: Vec = (0..40).step_by(2).collect(); + + acc.build_filter(FinalizeInput::Partitioned(routed_partitions( + &acc, + 4, + &build, + |partition, keys| { + if partition % 2 == 0 { + hash_map_of(keys) + } else { + in_list(keys) + } + }, + ))) .unwrap(); assert!( @@ -1383,7 +1452,9 @@ mod tests { "expected a merged lookup, got: {}", current_expr(&acc) ); - assert_matches(&acc, &[2, 7, 9, 3, 8], &[true, true, true, false, false]); + let probe: Vec = (0..40).collect(); + let expected: Vec = probe.iter().map(|key| key % 2 == 0).collect(); + assert_matches(&acc, &probe, &expected); } /// IN list partitions whose combined list passes the IN list limits merge @@ -1394,12 +1465,14 @@ mod tests { fn partitioned_in_lists_past_limit_merge_into_hash_lookup() { let mut acc = make_partitioned_expr_accumulator_for_test(3); acc.in_list_max_values = 2; + let build: Vec = (1..=9).collect(); - acc.build_filter(FinalizeInput::Partitioned(vec![ - reported(in_list(&[1]), no_bounds()), - reported(in_list(&[2]), no_bounds()), - reported(in_list(&[3]), no_bounds()), - ])) + acc.build_filter(FinalizeInput::Partitioned(routed_partitions( + &acc, + 3, + &build, + |_, keys| in_list(keys), + ))) .unwrap(); let expr = current_expr(&acc); @@ -1408,7 +1481,9 @@ mod tests { && expr.downcast_ref::().is_none(), "expected a hash lookup, got: {expr}" ); - assert_matches(&acc, &[1, 2, 3, 4], &[true, true, true, false]); + let probe: Vec = (0..20).collect(); + let expected: Vec = probe.iter().map(|key| (1..=9).contains(key)).collect(); + assert_matches(&acc, &probe, &expected); } /// Evaluates the current filter over probe keys `keys`. diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index a20cfdf4f5e05..7ddc381bf0d78 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -165,9 +165,6 @@ pub trait JoinHashMapType: Send + Sync { /// Returns the number of entries in the join hash map. fn len(&self) -> usize; - - /// Returns every distinct hash value stored in the map. - fn hash_values(&self) -> Box + '_>; } pub struct JoinHashMapU32 { @@ -254,10 +251,6 @@ impl JoinHashMapType for JoinHashMapU32 { fn len(&self) -> usize { self.map.len() } - - fn hash_values(&self) -> Box + '_> { - Box::new(self.map.iter().map(|&(hash, _)| hash)) - } } pub struct JoinHashMapU64 { @@ -344,10 +337,6 @@ impl JoinHashMapType for JoinHashMapU64 { fn len(&self) -> usize { self.map.len() } - - fn hash_values(&self) -> Box + '_> { - Box::new(self.map.iter().map(|&(hash, _)| hash)) - } } use crate::joins::chain::traverse_chain; @@ -367,78 +356,29 @@ where { let mut chained = false; for (row, &hash_value) in iter { - chained |= insert_hash(map, next, row, hash_value, deleted_offset); - } - chained -} - -/// Inserts `row` under `hash_value`, chaining it onto any row already stored -/// there. Returns `true` if it was chained. -#[inline] -fn insert_hash( - map: &mut HashTable<(u64, T)>, - next: &mut [T], - row: usize, - hash_value: u64, - deleted_offset: usize, -) -> bool -where - T: Copy + TryFrom, - >::Error: Debug, -{ - let entry = map.entry( - hash_value, - |&(hash, _)| hash_value == hash, - |&(hash, _)| hash, - ); - - match entry { - Occupied(mut occupied_entry) => { - // Already exists: add index to next array - let (_, index) = occupied_entry.get_mut(); - let prev_index = *index; - // Store new value inside hashmap - *index = T::try_from(row + 1).unwrap(); - // Update chained Vec at `row` with previous value - next[row - deleted_offset] = prev_index; - true - } - Vacant(vacant_entry) => { - vacant_entry.insert((hash_value, T::try_from(row + 1).unwrap())); - false - } - } -} + let entry = map.entry( + hash_value, + |&(hash, _)| hash_value == hash, + |&(hash, _)| hash, + ); -/// Builds one map holding the union of `maps`' hash values and `hashes`, so a -/// single [`JoinHashMapType::contain_hashes`] call checks membership in all of -/// them. -/// -/// Its row indices number the merged hash values rather than rows of any build -/// batch, so it is only meaningful for membership checks, never to produce -/// join matches. -pub(crate) fn union_hash_maps( - maps: &[&dyn JoinHashMapType], - hashes: &[u64], -) -> Box { - let len = maps.iter().map(|map| map.len()).sum::() + hashes.len(); - let hashes = maps - .iter() - .flat_map(|map| map.hash_values()) - .chain(hashes.iter().copied()); - if len > u32::MAX as usize { - let mut union = JoinHashMapU64::with_capacity(len); - for (row, hash) in hashes.enumerate() { - union.chained |= insert_hash(&mut union.map, &mut union.next, row, hash, 0); - } - Box::new(union) - } else { - let mut union = JoinHashMapU32::with_capacity(len); - for (row, hash) in hashes.enumerate() { - union.chained |= insert_hash(&mut union.map, &mut union.next, row, hash, 0); + match entry { + Occupied(mut occupied_entry) => { + // Already exists: add index to next array + let (_, index) = occupied_entry.get_mut(); + let prev_index = *index; + // Store new value inside hashmap + *index = T::try_from(row + 1).unwrap(); + // Update chained Vec at `row` with previous value + next[row - deleted_offset] = prev_index; + chained = true; + } + Vacant(vacant_entry) => { + vacant_entry.insert((hash_value, T::try_from(row + 1).unwrap())); + } } - Box::new(union) } + chained } pub fn get_matched_indices<'a, T>( diff --git a/datafusion/physical-plan/src/joins/stream_join_utils.rs b/datafusion/physical-plan/src/joins/stream_join_utils.rs index 0834b2b465622..90ce0d87a1732 100644 --- a/datafusion/physical-plan/src/joins/stream_join_utils.rs +++ b/datafusion/physical-plan/src/joins/stream_join_utils.rs @@ -117,10 +117,6 @@ impl JoinHashMapType for PruningJoinHashMap { fn len(&self) -> usize { self.map.len() } - - fn hash_values(&self) -> Box + '_> { - Box::new(self.map.iter().map(|&(hash, _)| hash)) - } } /// The `PruningJoinHashMap` is similar to a regular `JoinHashMap`, but with diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index b4984ab5b234f..91ed1c99501ae 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -796,6 +796,20 @@ impl RangeExpr { pub fn sort_options(&self) -> &[SortOptions] { self.router.sort_options() } + + /// Appends the partition of each row of `arrays`, the evaluated + /// `on_columns`, to `partition_ids`. + pub(crate) fn route_partition_ids( + &self, + arrays: &[arrow::array::ArrayRef], + partition_ids: &mut Vec, + ) -> Result<()> { + if self.router.split_points().is_empty() { + partition_ids.resize(partition_ids.len() + arrays[0].len(), 0); + return Ok(()); + } + self.router.route_partition_ids(arrays, partition_ids) + } } impl Display for RangeExpr {