feat(perf): compute routing and lookup hashes once

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}]) ]
     "
     );
 

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,
                         ))
                     })))
                 })
@@ -10124,19 +10125,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 <membership over every partition>`.
         assert!(
             filter.starts_with("c2@0 IS NULL OR"),
             "expected the IS NULL disjunct to wrap the whole routed filter, got: {filter}"

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<PhysicalExprRef>,
     /// Random state for hashing (with seeds preserved for serialization)
     random_state: SeededRandomState,
-    /// Map to check against (hash table or array map)
-    map: Arc<Map>,
+    /// 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<Arc<KeyRangeBitmap>>,
 }
+
+/// The build side a [`HashTableLookupExpr`] checks probe keys against.
+#[derive(Clone)]
+enum Lookup {
+    /// One map holding every build key.
+    Single(Arc<Map>),
+    /// 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<Arc<Map>>]>,
+    },
+}
+
+/// 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<RangeExpr>),
+}
+
+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<u64>,
+    ) -> 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<u64> {
+        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<PhysicalExprRef>,
+        random_state: SeededRandomState,
+        router: PartitionRouter,
+        maps: Vec<Option<Arc<Map>>>,
+        description: String,
+        pruning_bitmap: Option<Arc<KeyRangeBitmap>>,
+    ) -> Result<Self> {
+        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<KeyRangeBitmap>> {
         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<Arc<Map>>],
+    ) -> Result<BooleanArray> {
+        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::<Vec<_>>()
             .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;