Fix partitioned hash join dynamic filter routing

Author: Dandandan

datafusion/core/tests/physical_optimizer/filter_pushdown.rs

@@ -1083,7 +1083,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 [ CASE hash_repartition WHEN 2 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 8 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 ]
     "
     );
 

datafusion/expr/src/logical_plan/plan.rs

@@ -3773,6 +3773,8 @@ impl PartialOrd for Aggregate {
 /// Returns 0 when no grouping set is duplicated.
 fn max_grouping_set_duplicate_ordinal(group_expr: &[Expr]) -> usize {
     if let Some(Expr::GroupingSet(GroupingSet::GroupingSets(sets))) = group_expr.first() {
+        #[allow(clippy::allow_attributes, clippy::mutable_key_type)]
+        // Expr contains Arc with interior mutability but is intentionally used as hash key
         let mut counts: HashMap<&[Expr], usize> = HashMap::new();
         for set in sets {
             *counts.entry(set).or_insert(0) += 1;

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

@@ -33,6 +33,7 @@ use datafusion_physical_expr_common::physical_expr::{
 };
 
 use crate::joins::Map;
+use crate::repartition::hash_to_partition;
 
 /// RandomState wrapper that preserves the seed used to create it.
 ///
@@ -201,6 +202,127 @@ impl PhysicalExpr for HashExpr {
     }
 }
 
+/// Physical expression that computes the output partition id for hash repartitioning.
+///
+/// This expression must stay aligned with `RepartitionExec` because partitioned
+/// dynamic filters use it to route probe rows to the build-side hash table for
+/// the matching hash partition.
+pub(crate) struct HashPartitionExpr {
+    /// Columns to hash
+    on_columns: Vec<PhysicalExprRef>,
+    /// Random state for hashing (with seeds preserved for serialization)
+    random_state: SeededRandomState,
+    /// Number of output partitions
+    partition_count: usize,
+    /// Description for display
+    description: String,
+}
+
+impl HashPartitionExpr {
+    pub(crate) fn new(
+        on_columns: Vec<PhysicalExprRef>,
+        random_state: SeededRandomState,
+        partition_count: usize,
+        description: String,
+    ) -> Self {
+        Self {
+            on_columns,
+            random_state,
+            partition_count,
+            description,
+        }
+    }
+}
+
+impl std::fmt::Debug for HashPartitionExpr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let cols = self
+            .on_columns
+            .iter()
+            .map(|e| e.to_string())
+            .collect::<Vec<_>>()
+            .join(", ");
+        let seed = self.random_state.seed();
+        let partition_count = self.partition_count;
+        write!(
+            f,
+            "{}({cols}, [{seed}], partitions={partition_count})",
+            self.description
+        )
+    }
+}
+
+impl Hash for HashPartitionExpr {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.on_columns.dyn_hash(state);
+        self.description.hash(state);
+        self.random_state.seed().hash(state);
+        self.partition_count.hash(state);
+    }
+}
+
+impl PartialEq for HashPartitionExpr {
+    fn eq(&self, other: &Self) -> bool {
+        self.on_columns == other.on_columns
+            && self.description == other.description
+            && self.random_state.seed() == other.random_state.seed()
+            && self.partition_count == other.partition_count
+    }
+}
+
+impl Eq for HashPartitionExpr {}
+
+impl Display for HashPartitionExpr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.description)
+    }
+}
+
+impl PhysicalExpr for HashPartitionExpr {
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        self.on_columns.iter().collect()
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        Ok(Arc::new(HashPartitionExpr::new(
+            children,
+            self.random_state.clone(),
+            self.partition_count,
+            self.description.clone(),
+        )))
+    }
+
+    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+        Ok(DataType::UInt64)
+    }
+
+    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+        Ok(false)
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+        let keys_values = evaluate_columns(&self.on_columns, batch)?;
+
+        with_hashes(&keys_values, self.random_state.random_state(), |hashes| {
+            let partitions = hashes
+                .iter()
+                .map(|hash| hash_to_partition(*hash, self.partition_count) as u64)
+                .collect::<Vec<_>>();
+
+            Ok(ColumnarValue::Array(Arc::new(UInt64Array::from(
+                partitions,
+            ))))
+        })
+    }
+
+    fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.description)
+    }
+}
+
 /// Physical expression that checks join keys in a [`Map`] (hash table or array map).
 ///
 /// Returns a [`BooleanArray`](arrow::array::BooleanArray) indicating if join keys (from `on_columns`) exist in the map.

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

@@ -28,7 +28,7 @@ use crate::joins::PartitionMode;
 use crate::joins::hash_join::exec::HASH_JOIN_SEED;
 use crate::joins::hash_join::inlist_builder::build_struct_fields;
 use crate::joins::hash_join::partitioned_hash_eval::{
-    HashExpr, HashTableLookupExpr, SeededRandomState,
+    HashPartitionExpr, HashTableLookupExpr, SeededRandomState,
 };
 use arrow::array::ArrayRef;
 use arrow::datatypes::{DataType, Field, Schema};
@@ -249,8 +249,7 @@ pub(crate) struct SharedBuildAccumulator {
     dynamic_filter: Arc<DynamicFilterPhysicalExpr>,
     /// Right side join expressions needed for creating filter expressions
     on_right: Vec<PhysicalExprRef>,
-    /// Random state for partitioning (RepartitionExec's hash function with 0,0,0,0 seeds)
-    /// Used for PartitionedHashLookupPhysicalExpr
+    /// Random state for partitioning, matching RepartitionExec's hash function.
     repartition_random_state: SeededRandomState,
     /// Schema of the probe (right) side for evaluating filter expressions
     probe_schema: Arc<Schema>,
@@ -594,18 +593,13 @@ impl SharedBuildAccumulator {
             },
             FinalizeInput::Partitioned(partitions) => {
                 let num_partitions = partitions.len();
-                let routing_hash_expr = Arc::new(HashExpr::new(
+                let partition_expr = Arc::new(HashPartitionExpr::new(
                     self.on_right.clone(),
                     self.repartition_random_state.clone(),
+                    num_partitions,
                     "hash_repartition".to_string(),
                 )) as Arc<dyn PhysicalExpr>;
 
-                let modulo_expr = Arc::new(BinaryExpr::new(
-                    routing_hash_expr,
-                    Operator::Modulo,
-                    lit(ScalarValue::UInt64(Some(num_partitions as u64))),
-                )) as Arc<dyn PhysicalExpr>;
-
                 let mut real_branches = Vec::new();
                 let mut empty_partition_ids = Vec::new();
                 let mut has_canceled_unknown = false;
@@ -665,7 +659,7 @@ impl SharedBuildAccumulator {
                         lit(true)
                     } else {
                         Arc::new(CaseExpr::try_new(
-                            Some(modulo_expr),
+                            Some(Arc::clone(&partition_expr)),
                             when_then_branches,
                             Some(lit(true)),
                         )?) as Arc<dyn PhysicalExpr>
@@ -678,7 +672,7 @@ impl SharedBuildAccumulator {
                     Arc::clone(&real_branches[0].1)
                 } else {
                     Arc::new(CaseExpr::try_new(
-                        Some(modulo_expr),
+                        Some(partition_expr),
                         real_branches,
                         Some(lit(false)),
                     )?) as Arc<dyn PhysicalExpr>

datafusion/physical-plan/src/repartition/mod.rs

@@ -438,6 +438,10 @@ enum BatchPartitionerState {
 /// executions and runs.
 pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_seed(0);
 
+pub(crate) fn hash_to_partition(hash: u64, partitions: usize) -> usize {
+    (((hash as u128) * (partitions as u128)) >> 64) as usize
+}
+
 impl BatchPartitioner {
     /// Create a new [`BatchPartitioner`] for hash-based repartitioning.
     ///
@@ -597,8 +601,7 @@ impl BatchPartitioner {
                     indices.iter_mut().for_each(|v| v.clear());
 
                     for (index, hash) in hash_buffer.iter().enumerate() {
-                        let part =
-                            (((*hash as u128) * (*partitions as u128)) >> 64) as usize;
+                        let part = hash_to_partition(*hash, *partitions);
                         indices[part].push(index as u32);
                     }
 

datafusion/sqllogictest/test_files/clickbench_extended.slt

@@ -45,8 +45,8 @@ SELECT "BrowserCountry",  COUNT(DISTINCT "SocialNetwork"), COUNT(DISTINCT "HitCo
 query IIIRRRR
 SELECT "SocialSourceNetworkID", "RegionID", COUNT(*), AVG("Age"), AVG("ParamPrice"), STDDEV("ParamPrice") as s, VAR("ParamPrice")  FROM hits GROUP BY "SocialSourceNetworkID", "RegionID" HAVING s IS NOT NULL ORDER BY s DESC LIMIT 10;
 ----
-0 839 6 0 0 0 0
 0 197 2 0 0 0 0
+0 839 6 0 0 0 0
 
 query IIIIII
 SELECT "ClientIP", "WatchID",  COUNT(*) c, MIN("ResponseStartTiming") tmin, MEDIAN("ResponseStartTiming") tmed, MAX("ResponseStartTiming") tmax FROM hits WHERE "JavaEnable" = 0  GROUP BY  "ClientIP", "WatchID" HAVING c > 1 ORDER BY tmed DESC LIMIT 10;