feat(physical-plan): Allow co-partitioned Partitioning::Range inputs for left-side hash joins

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -933,6 +933,79 @@ fn range_window_rehashes_incompatible_range_partitioning() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn range_left_mark_hash_join_reuses_range_partitioning() -> Result<()> {
+    let left = parquet_exec_with_output_partitioning(range_partitioning(
+        "a",
+        [10, 20, 30],
+        SortOptions::default(),
+    )?);
+    let right = parquet_exec_with_output_partitioning(range_partitioning(
+        "a",
+        [10, 20, 30],
+        SortOptions::default(),
+    )?);
+    let join_on = vec![(
+        Arc::new(Column::new_with_schema("a", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("a", &right.schema())?) as _,
+    )];
+    let join = hash_join_exec(left, right, &join_on, &JoinType::LeftMark);
+
+    let plan = TestConfig::default()
+        .with_query_execution_partitions(4)
+        .to_plan(join, &DISTRIB_DISTRIB_SORT);
+
+    assert_plan!(
+        plan,
+        @r"
+    HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(a@0, a@0)]
+      DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+      DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+    "
+    );
+
+    Ok(())
+}
+
+#[test]
+fn range_left_anti_hash_join_rehashes_incompatible_sort_options() -> Result<()> {
+    let left = parquet_exec_with_output_partitioning(range_partitioning(
+        "a",
+        [10, 20, 30],
+        SortOptions::default(),
+    )?);
+    let right = parquet_exec_with_output_partitioning(range_partitioning(
+        "a",
+        [30, 20, 10],
+        SortOptions {
+            descending: true,
+            nulls_first: true,
+        },
+    )?);
+    let join_on = vec![(
+        Arc::new(Column::new_with_schema("a", &left.schema())?) as _,
+        Arc::new(Column::new_with_schema("a", &right.schema())?) as _,
+    )];
+    let join = hash_join_exec(left, right, &join_on, &JoinType::LeftAnti);
+
+    let plan = TestConfig::default()
+        .with_query_execution_partitions(4)
+        .to_plan(join, &DISTRIB_DISTRIB_SORT);
+
+    assert_plan!(
+        plan,
+        @r"
+    HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(a@0, a@0)]
+      RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4
+        DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+      RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4
+        DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 DESC], [(30), (20), (10)], 4), file_type=parquet
+    "
+    );
+
+    Ok(())
+}
+
 #[test]
 fn multi_hash_joins() -> Result<()> {
     let left = parquet_exec();

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

@@ -1293,7 +1293,16 @@ impl ExecutionPlan for HashJoinExec {
             ]),
         };
 
-        if self.mode == PartitionMode::Partitioned && self.join_type == JoinType::Inner {
+        if self.mode == PartitionMode::Partitioned
+            && matches!(
+                self.join_type,
+                JoinType::Inner
+                    | JoinType::Left
+                    | JoinType::LeftSemi
+                    | JoinType::LeftAnti
+                    | JoinType::LeftMark
+            )
+        {
             requirements.allow_range_satisfaction_for_key_partitioning()
         } else {
             requirements

datafusion/sqllogictest/src/test_context/range_partitioning.rs

@@ -121,7 +121,7 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
         "range_partitioned_shifted",
         Path::new(env!("CARGO_MANIFEST_DIR"))
             .join("test_files/scratch_range_partitioning/range_partitioned_shifted"),
-        schema,
+        Arc::clone(&schema),
         [
             "1,1,10\n5,2,50\n10,1,100\n",
             "15,2,150\n",
@@ -130,6 +130,34 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
         ],
         Some(shifted_output_partitioning),
     );
+
+    // Same rows as `range_partitioned` but split into only three range
+    // partitions on `range_key`. Used to exercise the co-partition check when
+    // two Range inputs disagree on partition count.
+    let narrow_output_partitioning = Partitioning::Range(
+        RangePartitioning::try_new(
+            vec![col("range_key").sort(true, true)],
+            vec![
+                SplitPoint::new(vec![ScalarValue::Int32(Some(10))]),
+                SplitPoint::new(vec![ScalarValue::Int32(Some(20))]),
+            ],
+        )
+        .expect("range partitioning should be valid"),
+    );
+
+    register_csv_listing_table(
+        ctx,
+        "range_partitioned_narrow",
+        Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("test_files/scratch_range_partitioning/range_partitioned_narrow"),
+        schema,
+        [
+            "1,1,10\n5,2,50\n",
+            "10,1,100\n15,2,150\n",
+            "20,1,200\n25,2,250\n30,1,300\n35,2,350\n",
+        ],
+        Some(narrow_output_partitioning),
+    );
 }
 
 fn register_csv_listing_table(

datafusion/sqllogictest/test_files/range_partitioning.slt

@@ -395,27 +395,123 @@ ORDER BY l.non_range_key, l.value, r.value;
 2 350 350
 
 ##########
-# TEST 12: Non-Inner Range Join Repartitions
-# Only inner partitioned hash joins opt in to Range satisfying KeyPartitioned
-# requirements. Non-inner joins keep using Hash repartitioning.
+# TEST 12: Left-Side Range Hash Joins
+# Left-side partitioned hash joins also opt in to Range satisfying
+# KeyPartitioned requirements. Compatible Range layouts satisfy both the
+# per-child key requirements and the cross-child layout requirement, so no
+# Hash repartitioning is inserted. Incompatible split points, partition counts,
+# or join key expressions still use Hash repartitioning to repair the inputs.
 ##########
 
 query TT
 EXPLAIN SELECT l.range_key, l.value, r.value
 FROM range_partitioned l
-LEFT JOIN range_partitioned r ON l.range_key = r.range_key;
+LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r
+ON l.range_key = r.range_key;
+----
+physical_plan
+01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3]
+02)--DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+03)--FilterExec: value@1 <= 150
+04)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+
+query III
+SELECT l.range_key, l.value, r.value
+FROM range_partitioned l
+LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r
+ON l.range_key = r.range_key
+ORDER BY l.range_key;
+----
+1 10 10
+5 50 50
+10 100 100
+15 150 150
+20 200 NULL
+25 250 NULL
+30 300 NULL
+35 350 NULL
+
+query TT
+EXPLAIN SELECT l.range_key, l.value
+FROM range_partitioned l
+LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r
+ON l.range_key = r.range_key;
+----
+physical_plan
+01)HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(range_key@0, range_key@0)]
+02)--DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+03)--FilterExec: value@1 <= 150, projection=[range_key@0]
+04)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+
+query II
+SELECT l.range_key, l.value
+FROM range_partitioned l
+LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r
+ON l.range_key = r.range_key
+ORDER BY l.range_key;
+----
+1 10
+5 50
+10 100
+15 150
+
+query TT
+EXPLAIN SELECT l.range_key, l.value
+FROM range_partitioned l
+LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r
+ON l.range_key = r.range_key;
+----
+physical_plan
+01)HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(range_key@0, range_key@0)]
+02)--DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+03)--FilterExec: value@1 <= 150, projection=[range_key@0]
+04)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+
+query II
+SELECT l.range_key, l.value
+FROM range_partitioned l
+LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r
+ON l.range_key = r.range_key
+ORDER BY l.range_key;
+----
+20 200
+25 250
+30 300
+35 350
+
+# Range([range_key]) is only a subset of the composite join key, so the
+# co-partitioned hash join requirement is repaired with Hash repartitioning.
+query TT
+EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value
+FROM range_partitioned l
+LEFT JOIN (SELECT range_key, non_range_key, value FROM range_partitioned WHERE value <= 150) r
+ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key;
+----
+physical_plan
+01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0), (non_range_key@1, non_range_key@1)], projection=[range_key@0, non_range_key@1, value@2, value@5]
+02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4
+03)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+04)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4
+05)----FilterExec: value@2 <= 150
+06)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+
+# Different split points do not satisfy the co-partitioned layout requirement.
+query TT
+EXPLAIN SELECT l.range_key, l.value, r.value
+FROM range_partitioned l
+LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key;
 ----
 physical_plan
 01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3]
 02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4
 03)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
 04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4
-05)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+05)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false
 
 query III
 SELECT l.range_key, l.value, r.value
 FROM range_partitioned l
-LEFT JOIN range_partitioned r ON l.range_key = r.range_key
+LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key
 ORDER BY l.range_key;
 ----
 1 10 10
@@ -427,6 +523,77 @@ ORDER BY l.range_key;
 30 300 300
 35 350 350
 
+# Different partition counts do not satisfy the co-partitioned layout
+# requirement.
+query TT
+EXPLAIN SELECT l.range_key, l.value, r.value
+FROM range_partitioned l
+LEFT JOIN range_partitioned_narrow r ON l.range_key = r.range_key;
+----
+physical_plan
+01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3]
+02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4
+03)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=3
+05)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20)], 3), file_type=csv, has_header=false
+
+# Range([range_key]) does not satisfy a join keyed on non_range_key.
+query TT
+EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value
+FROM range_partitioned l
+LEFT JOIN range_partitioned r ON l.non_range_key = r.non_range_key;
+----
+physical_plan
+01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(non_range_key@1, non_range_key@0)], projection=[range_key@0, non_range_key@1, value@2, value@4]
+02)--RepartitionExec: partitioning=Hash([non_range_key@1], 4), input_partitions=4
+03)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+04)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4
+05)----DataSourceExec: file_groups=<slt:ignore>, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false
+
+# SQL IN subqueries decorrelate to LeftMark joins. These queries pin matched,
+# unmatched, and NULL marker behavior over compatible Range inputs.
+query TT
+EXPLAIN SELECT l.range_key, l.value
+FROM range_partitioned l
+WHERE l.non_range_key = 2 OR l.range_key IN (
+  SELECT range_key FROM range_partitioned WHERE value <= 150);
+----
+physical_plan
+01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2]
+02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)]
+03)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+04)----FilterExec: value@1 <= 150, projection=[range_key@0]
+05)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
+
+query II
+SELECT l.range_key, l.value
+FROM range_partitioned l
+WHERE l.non_range_key = 2 OR l.range_key IN (
+  SELECT range_key FROM range_partitioned WHERE value <= 150)
+ORDER BY l.range_key;
+----
+1 10
+5 50
+10 100
+15 150
+25 250
+35 350
+
+query II
+SELECT l.range_key, l.value
+FROM range_partitioned l
+WHERE l.non_range_key = 2 OR l.range_key IN (
+  SELECT CASE WHEN value <= 150 THEN range_key ELSE NULL END
+  FROM range_partitioned)
+ORDER BY l.range_key;
+----
+1 10
+5 50
+10 100
+15 150
+25 250
+35 350
+
 ##########
 # TEST 13: Compatible Range Join Repartitions to Increase Parallelism
 # Co-partitioning satisfaction does not prevent a repartition that increases