Enable dynamic filters for range-partitioned joins

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

@@ -886,9 +886,6 @@ impl HashJoinExec {
         if self.mode == PartitionMode::Partitioned
             && !self.has_partitioned_dynamic_filter_routing()
         {
-            // TODO: support partition-routed dynamic filters for compatible
-            // range co-partitioned joins.
-            // <https://github.com/apache/datafusion/issues/23376>.
             return false;
         }
 
@@ -904,6 +901,14 @@ impl HashJoinExec {
                 Partitioning::Hash(_, left_partition_count),
                 Partitioning::Hash(_, right_partition_count),
             ) => left_partition_count == right_partition_count,
+            (Partitioning::Range(_), Partitioning::Range(_)) => {
+                let children = [self.left.as_ref(), self.right.as_ref()];
+                matches!(
+                    self.input_distribution_requirements()
+                        .unsatisfied_co_partitioned_children(self.name(), &children),
+                    Ok(unsatisfied) if unsatisfied.is_empty()
+                )
+            }
             (left_partitioning, right_partitioning) => {
                 left_partitioning.partition_count() == 1
                     && right_partitioning.partition_count() == 1
@@ -6759,8 +6764,7 @@ mod tests {
     }
 
     #[test]
-    fn test_partitioned_dynamic_filter_pushdown_rejects_range_partitioning() -> Result<()>
-    {
+    fn test_partitioned_dynamic_filter_pushdown_range_partitioning() -> Result<()> {
         let (left_schema, right_schema, on) = build_schema_and_on()?;
         let left_partitioning = Partitioning::Range(RangePartitioning::try_new(
             [PhysicalSortExpr {
@@ -6783,7 +6787,7 @@ mod tests {
             left_partitioning,
         )?);
         let right = Arc::new(PartitionedTestExec::try_new(
-            right_schema,
+            Arc::clone(&right_schema),
             right_partitioning,
         )?);
 
@@ -6794,8 +6798,35 @@ mod tests {
             .enable_join_dynamic_filter_pushdown = true;
 
         let join = HashJoinExec::try_new(
-            left,
+            Arc::clone(&left) as Arc<dyn ExecutionPlan>,
             right,
+            on.clone(),
+            None,
+            &JoinType::Inner,
+            None,
+            PartitionMode::Partitioned,
+            NullEquality::NullEqualsNothing,
+            false,
+        )?;
+
+        assert!(join.allow_join_dynamic_filter_pushdown(session_config.options()));
+
+        let mismatched_right_partitioning =
+            Partitioning::Range(RangePartitioning::try_new(
+                [PhysicalSortExpr {
+                    expr: Arc::clone(&on[0].1),
+                    options: Default::default(),
+                }]
+                .into(),
+                vec![SplitPoint::new(vec![ScalarValue::Int32(Some(11))])],
+            )?);
+        let mismatched_right = Arc::new(PartitionedTestExec::try_new(
+            right_schema,
+            mismatched_right_partitioning,
+        )?);
+        let mismatched_join = HashJoinExec::try_new(
+            left,
+            mismatched_right,
             on,
             None,
             &JoinType::Inner,
@@ -6805,6 +6836,67 @@ mod tests {
             false,
         )?;
 
+        assert!(
+            !mismatched_join.allow_join_dynamic_filter_pushdown(session_config.options())
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_partitioned_dynamic_filter_pushdown_rejects_float_zero_split() -> Result<()> {
+        let left_schema = Arc::new(Schema::new(vec![Field::new(
+            "left_key",
+            DataType::Float64,
+            false,
+        )]));
+        let right_schema = Arc::new(Schema::new(vec![Field::new(
+            "right_key",
+            DataType::Float64,
+            false,
+        )]));
+        let left_key = Arc::new(Column::new("left_key", 0)) as PhysicalExprRef;
+        let right_key = Arc::new(Column::new("right_key", 0)) as PhysicalExprRef;
+        let split_points = vec![SplitPoint::new(vec![ScalarValue::Float64(Some(0.0))])];
+        let left = Arc::new(PartitionedTestExec::try_new(
+            left_schema,
+            Partitioning::Range(RangePartitioning::try_new(
+                [PhysicalSortExpr::new(
+                    Arc::clone(&left_key),
+                    Default::default(),
+                )]
+                .into(),
+                split_points.clone(),
+            )?),
+        )?);
+        let right = Arc::new(PartitionedTestExec::try_new(
+            right_schema,
+            Partitioning::Range(RangePartitioning::try_new(
+                [PhysicalSortExpr::new(
+                    Arc::clone(&right_key),
+                    Default::default(),
+                )]
+                .into(),
+                split_points,
+            )?),
+        )?);
+        let join = HashJoinExec::try_new(
+            left,
+            right,
+            vec![(left_key, right_key)],
+            None,
+            &JoinType::Inner,
+            None,
+            PartitionMode::Partitioned,
+            NullEquality::NullEqualsNothing,
+            false,
+        )?;
+        let mut session_config = SessionConfig::default();
+        session_config
+            .options_mut()
+            .optimizer
+            .enable_join_dynamic_filter_pushdown = true;
+
         assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options()));
 
         Ok(())