Enforce co-partitioning for sort merge and symmetric hash joins

datafusion/core/tests/physical_optimizer/sanity_checker.rs

@@ -20,21 +20,23 @@ use std::sync::Arc;
 
 use crate::physical_optimizer::test_utils::{
     bounded_window_exec, global_limit_exec, hash_join_exec, local_limit_exec,
-    memory_exec, projection_exec, repartition_exec, sort_exec, sort_expr,
-    sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
+    memory_exec, projection_exec, repartition_exec, sort_exec,
+    sort_exec_with_preserve_partitioning, sort_expr, sort_expr_options,
+    sort_merge_join_exec, sort_preserving_merge_exec, union_exec,
 };
 
 use arrow::compute::SortOptions;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable};
 use datafusion::prelude::{CsvReadOptions, SessionContext};
 use datafusion_common::config::ConfigOptions;
-use datafusion_common::{JoinType, Result, ScalarValue};
+use datafusion_common::{JoinType, NullEquality, Result, ScalarValue};
 use datafusion_physical_expr::expressions::{Literal, col};
 use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan;
+use datafusion_physical_plan::joins::{StreamJoinPartitionMode, SymmetricHashJoinExec};
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::{ExecutionPlan, displayable};
 
@@ -417,6 +419,26 @@ fn range_partitioned_exec(
         .map(|exec| Arc::new(exec) as Arc<dyn ExecutionPlan>)
 }
 
+fn hash_partitioned_exec(
+    schema: &SchemaRef,
+    key: &str,
+    partition_count: usize,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let partitioning = Partitioning::Hash(vec![col(key, schema)?], partition_count);
+    RepartitionExec::try_new(memory_exec(schema), partitioning)
+        .map(|exec| Arc::new(exec) as Arc<dyn ExecutionPlan>)
+}
+
+fn sorted_hash_partitioned_exec(
+    schema: &SchemaRef,
+    key: &str,
+    partition_count: usize,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    let input = hash_partitioned_exec(schema, key, partition_count)?;
+    let ordering: LexOrdering = [sort_expr(key, schema)].into();
+    Ok(sort_exec_with_preserve_partitioning(ordering, input))
+}
+
 #[test]
 fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> {
     let schema = create_test_schema2();
@@ -443,6 +465,64 @@ fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_sort_merge_join_requires_co_partitioned_children() -> Result<()> {
+    let schema = create_test_schema2();
+    let join_on = vec![(col("a", &schema)?, col("a", &schema)?)];
+
+    let compatible_join = sort_merge_join_exec(
+        sorted_hash_partitioned_exec(&schema, "a", 4)?,
+        sorted_hash_partitioned_exec(&schema, "a", 4)?,
+        &join_on,
+        &JoinType::Inner,
+    );
+    assert_sanity_check(&compatible_join, true);
+
+    let incompatible_join = sort_merge_join_exec(
+        sorted_hash_partitioned_exec(&schema, "a", 4)?,
+        sorted_hash_partitioned_exec(&schema, "a", 5)?,
+        &join_on,
+        &JoinType::Inner,
+    );
+    assert_sanity_check(&incompatible_join, false);
+
+    Ok(())
+}
+
+#[test]
+fn test_symmetric_hash_join_requires_co_partitioned_children() -> Result<()> {
+    let schema = create_test_schema2();
+    let join_on = vec![(col("a", &schema)?, col("a", &schema)?)];
+
+    let compatible_join = Arc::new(SymmetricHashJoinExec::try_new(
+        hash_partitioned_exec(&schema, "a", 4)?,
+        hash_partitioned_exec(&schema, "a", 4)?,
+        join_on.clone(),
+        None,
+        &JoinType::Inner,
+        NullEquality::NullEqualsNothing,
+        None,
+        None,
+        StreamJoinPartitionMode::Partitioned,
+    )?) as Arc<dyn ExecutionPlan>;
+    assert_sanity_check(&compatible_join, true);
+
+    let incompatible_join = Arc::new(SymmetricHashJoinExec::try_new(
+        hash_partitioned_exec(&schema, "a", 4)?,
+        hash_partitioned_exec(&schema, "a", 5)?,
+        join_on,
+        None,
+        &JoinType::Inner,
+        NullEquality::NullEqualsNothing,
+        None,
+        None,
+        StreamJoinPartitionMode::Partitioned,
+    )?) as Arc<dyn ExecutionPlan>;
+    assert_sanity_check(&incompatible_join, false);
+
+    Ok(())
+}
+
 #[tokio::test]
 /// Tests that plan is valid when the sort requirements are satisfied.
 async fn test_bounded_window_agg_sort_requirement() -> Result<()> {

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

@@ -41,7 +41,8 @@ use crate::spill::spill_manager::SpillManager;
 use crate::statistics::StatisticsArgs;
 use crate::{
     DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
-    PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties,
+    InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, Statistics,
+    check_if_same_properties,
 };
 
 use arrow::compute::SortOptions;
@@ -413,13 +414,13 @@ impl ExecutionPlan for SortMergeJoinExec {
         self.input_distribution_requirements().into_per_child()
     }
 
-    fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
         let (left_expr, right_expr) = self
             .on
             .iter()
             .map(|(l, r)| (Arc::clone(l), Arc::clone(r)))
             .unzip();
-        crate::InputDistributionRequirements::new(vec![
+        InputDistributionRequirements::co_partitioned(vec![
             Distribution::KeyPartitioned(left_expr),
             Distribution::KeyPartitioned(right_expr),
         ])

datafusion/physical-plan/src/joins/symmetric_hash_join.rs

@@ -53,7 +53,8 @@ use crate::projection::{
 use crate::stream::EmptyRecordBatchStream;
 use crate::{
     DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
-    PlanProperties, RecordBatchStream, SendableRecordBatchStream,
+    InputDistributionRequirements, PlanProperties, RecordBatchStream,
+    SendableRecordBatchStream,
     joins::StreamJoinPartitionMode,
     metrics::{ExecutionPlanMetricsSet, MetricsSet},
 };
@@ -415,23 +416,26 @@ impl ExecutionPlan for SymmetricHashJoinExec {
         self.input_distribution_requirements().into_per_child()
     }
 
-    fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
-        crate::InputDistributionRequirements::new(match self.mode {
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        match self.mode {
             StreamJoinPartitionMode::Partitioned => {
                 let (left_expr, right_expr) = self
                     .on
                     .iter()
                     .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _))
                     .unzip();
-                vec![
+                InputDistributionRequirements::co_partitioned(vec![
                     Distribution::KeyPartitioned(left_expr),
                     Distribution::KeyPartitioned(right_expr),
-                ]
+                ])
             }
             StreamJoinPartitionMode::SinglePartition => {
-                vec![Distribution::SinglePartition, Distribution::SinglePartition]
+                InputDistributionRequirements::new(vec![
+                    Distribution::SinglePartition,
+                    Distribution::SinglePartition,
+                ])
             }
-        })
+        }
     }
 
     fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {

datafusion/sqllogictest/src/test_context/range_partitioning.rs

@@ -25,7 +25,8 @@ use datafusion::datasource::file_format::csv::CsvFormat;
 use datafusion::datasource::listing::{
     ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
 };
-use datafusion::logical_expr::{Partitioning, RangePartitioning, col};
+use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable};
+use datafusion::logical_expr::{Partitioning, RangePartitioning, SortExpr, col};
 use datafusion::prelude::SessionContext;
 
 // ==============================================================================
@@ -52,11 +53,13 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
         .expect("range partitioning should be valid"),
     );
 
+    let range_table_dir = Path::new(env!("CARGO_MANIFEST_DIR"))
+        .join("test_files/scratch_range_partitioning/range_partitioned");
+
     register_csv_listing_table(
         ctx,
         "range_partitioned",
-        Path::new(env!("CARGO_MANIFEST_DIR"))
-            .join("test_files/scratch_range_partitioning/range_partitioned"),
+        &range_table_dir,
         Arc::clone(&schema),
         [
             "1,1,10\n5,2,50\n",
@@ -67,6 +70,14 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) {
         Some(output_partitioning),
     );
 
+    register_unbounded_stream_table(
+        ctx,
+        "unbounded_range_like",
+        range_table_dir.join("part-0.csv"),
+        Arc::clone(&schema),
+        vec![vec![col("range_key").sort(true, false)]],
+    );
+
     let shifted_output_partitioning = Partitioning::Range(
         RangePartitioning::try_new(
             vec![col("range_key").sort(true, true)],
@@ -133,3 +144,17 @@ fn register_csv_listing_table(
     ctx.register_table(name, Arc::new(table))
         .expect("test listing table registration should succeed");
 }
+
+fn register_unbounded_stream_table(
+    ctx: &SessionContext,
+    name: &str,
+    file_path: impl Into<std::path::PathBuf>,
+    schema: Arc<Schema>,
+    order: Vec<Vec<SortExpr>>,
+) {
+    let source = FileStreamProvider::new_file(schema, file_path.into());
+    let config = StreamConfig::new(Arc::new(source)).with_order(order);
+
+    ctx.register_table(name, Arc::new(StreamTable::new(Arc::new(config))))
+        .expect("test stream table registration should succeed");
+}

datafusion/sqllogictest/test_files/range_partitioning.slt

@@ -634,6 +634,80 @@ ORDER BY l.range_key;
 30 600
 35 700
 
+##########
+# TEST 18: Sort Merge Join Repartitions Compatible Range Inputs
+# SortMergeJoinExec does not opt into Range satisfying KeyPartitioned. These
+# compatible Range inputs receive Hash repartitions, a future
+# Range satisfaction implementation should remove them.
+##########
+
+statement ok
+set datafusion.optimizer.prefer_hash_join = false;
+
+query TT
+EXPLAIN SELECT l.range_key, l.value, r.value
+FROM range_partitioned l
+JOIN range_partitioned r ON l.range_key = r.range_key;
+----
+physical_plan
+01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value]
+02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)]
+03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true]
+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
+06)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true]
+07)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4
+08)--------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
+JOIN range_partitioned 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 200
+25 250 250
+30 300 300
+35 350 350
+
+statement ok
+reset datafusion.optimizer.prefer_hash_join;
+
+##########
+# TEST 19: Symmetric Hash Join Repartitions Inputs
+# SymmetricHashJoinExec does not opt into Range satisfying KeyPartitioned.
+# The unbounded stream fixture does not yet expose Range partitioning,
+# so the join still receives Hash repartitions.
+##########
+
+statement ok
+set datafusion.optimizer.prefer_hash_join = true;
+
+query TT
+EXPLAIN SELECT l.range_key, l.value, r.value
+FROM unbounded_range_like l
+FULL JOIN unbounded_range_like r ON l.range_key = r.range_key;
+----
+physical_plan
+01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value]
+02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)]
+03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=1, maintains_sort_order=true
+04)------StreamingTableExec: partition_sizes=1, projection=[range_key, value], infinite_source=true, output_ordering=[range_key@0 ASC NULLS LAST]
+05)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=1, maintains_sort_order=true
+06)------StreamingTableExec: partition_sizes=1, projection=[range_key, value], infinite_source=true, output_ordering=[range_key@0 ASC NULLS LAST]
+
+query III rowsort
+SELECT l.range_key, l.value, r.value
+FROM unbounded_range_like l
+FULL JOIN unbounded_range_like r ON l.range_key = r.range_key;
+----
+1 10 10
+5 50 50
+
 statement ok
 reset datafusion.optimizer.prefer_hash_join;
 
@@ -644,7 +718,7 @@ statement ok
 reset datafusion.optimizer.preserve_file_partitions;
 
 ##########
-# TEST 18: Union of Range Partitioned Inputs
+# TEST 20: Union of Range Partitioned Inputs
 # Each input exposes Range partitioning on range_key. These changes do not add a
 # cross-child Range relationship for UNION ALL.
 ##########