Enforce co-partitioning for sort merge and symmetric hash joins

datafusion/catalog/src/streaming.rs

@@ -24,7 +24,10 @@ use async_trait::async_trait;
 use datafusion_common::{DFSchema, Result, plan_err};
 use datafusion_expr::{Expr, SortExpr, TableType};
 use datafusion_physical_expr::equivalence::project_ordering;
-use datafusion_physical_expr::{LexOrdering, create_physical_sort_exprs};
+use datafusion_physical_expr::projection::ProjectionMapping;
+use datafusion_physical_expr::{
+    EquivalenceProperties, LexOrdering, Partitioning, create_physical_sort_exprs,
+};
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec};
 use log::debug;
@@ -38,6 +41,7 @@ pub struct StreamingTable {
     partitions: Vec<Arc<dyn PartitionStream>>,
     infinite: bool,
     sort_order: Vec<SortExpr>,
+    output_partitioning: Option<Partitioning>,
 }
 
 impl StreamingTable {
@@ -62,6 +66,7 @@ impl StreamingTable {
             partitions,
             infinite: false,
             sort_order: vec![],
+            output_partitioning: None,
         })
     }
 
@@ -76,6 +81,33 @@ impl StreamingTable {
         self.sort_order = sort_order;
         self
     }
+
+    /// Declares the output partitioning of this streaming table.
+    ///
+    /// The partitioning expressions refer to the table schema before scan
+    /// projection. If a scan projection removes a partitioning expression, the
+    /// physical plan reports unknown partitioning.
+    pub fn with_output_partitioning(mut self, output_partitioning: Partitioning) -> Self {
+        self.output_partitioning = Some(output_partitioning);
+        self
+    }
+
+    fn output_partitioning(
+        &self,
+        projection: Option<&Vec<usize>>,
+    ) -> Result<Partitioning> {
+        let Some(output_partitioning) = &self.output_partitioning else {
+            return Ok(Partitioning::UnknownPartitioning(self.partitions.len()));
+        };
+        let Some(projection) = projection else {
+            return Ok(output_partitioning.clone());
+        };
+
+        let projection_mapping =
+            ProjectionMapping::from_indices(projection, &self.schema)?;
+        let eq_properties = EquivalenceProperties::new(Arc::clone(&self.schema));
+        Ok(output_partitioning.project(&projection_mapping, &eq_properties))
+    }
 }
 
 #[async_trait]
@@ -119,13 +151,16 @@ impl TableProvider for StreamingTable {
             vec![]
         };
 
-        Ok(Arc::new(StreamingTableExec::try_new(
+        let exec = StreamingTableExec::try_new(
             Arc::clone(&self.schema),
             self.partitions.clone(),
             projection,
             LexOrdering::new(physical_sort),
             self.infinite,
             limit,
-        )?))
+        )?
+        .with_output_partitioning(self.output_partitioning(projection)?)?;
+
+        Ok(Arc::new(exec))
     }
 }

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};
 
@@ -443,6 +445,77 @@ 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 ordering: LexOrdering = [sort_expr("a", &schema)].into();
+
+    let compatible_join = sort_merge_join_exec(
+        sort_exec_with_preserve_partitioning(
+            ordering.clone(),
+            range_partitioned_exec(&schema, "a", [10])?,
+        ),
+        sort_exec_with_preserve_partitioning(
+            ordering.clone(),
+            range_partitioned_exec(&schema, "a", [10])?,
+        ),
+        &join_on,
+        &JoinType::Inner,
+    );
+    assert_sanity_check(&compatible_join, true);
+
+    let incompatible_join = sort_merge_join_exec(
+        sort_exec_with_preserve_partitioning(
+            ordering.clone(),
+            range_partitioned_exec(&schema, "a", [10])?,
+        ),
+        sort_exec_with_preserve_partitioning(
+            ordering,
+            range_partitioned_exec(&schema, "a", [20])?,
+        ),
+        &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(
+        range_partitioned_exec(&schema, "a", [10])?,
+        range_partitioned_exec(&schema, "a", [10])?,
+        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(
+        range_partitioned_exec(&schema, "a", [10])?,
+        range_partitioned_exec(&schema, "a", [20])?,
+        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::{ChildStats, 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,16 +414,17 @@ 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),
         ])
+        .allow_range_satisfaction_for_key_partitioning()
     }
 
     fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {

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,27 @@ 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),
-                ]
+                ])
+                .allow_range_satisfaction_for_key_partitioning()
             }
             StreamJoinPartitionMode::SinglePartition => {
-                vec![Distribution::SinglePartition, Distribution::SinglePartition]
+                InputDistributionRequirements::new(vec![
+                    Distribution::SinglePartition,
+                    Distribution::SinglePartition,
+                ])
             }
-        })
+        }
     }
 
     fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {

datafusion/physical-plan/src/streaming.rs

@@ -35,6 +35,7 @@ use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream};
 use arrow::datatypes::{Schema, SchemaRef};
 use datafusion_common::{Result, internal_err, plan_err};
 use datafusion_execution::TaskContext;
+use datafusion_physical_expr::projection::ProjectionMapping;
 use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
 
 use async_trait::async_trait;
@@ -100,7 +101,7 @@ impl StreamingTableExec {
         let cache = Self::compute_properties(
             Arc::clone(&projected_schema),
             projected_output_ordering.clone(),
-            &partitions,
+            Partitioning::UnknownPartitioning(partitions.len()),
             infinite,
         );
         Ok(Self {
@@ -115,6 +116,25 @@ impl StreamingTableExec {
         })
     }
 
+    /// Declares the output partitioning of this stream.
+    ///
+    /// `output_partitioning` must describe this plan's current output and have
+    /// the same number of partitions as the stream.
+    pub fn with_output_partitioning(
+        mut self,
+        output_partitioning: Partitioning,
+    ) -> Result<Self> {
+        if output_partitioning.partition_count() != self.partitions.len() {
+            return plan_err!(
+                "Output partitioning has {} partitions but stream has {} partitions",
+                output_partitioning.partition_count(),
+                self.partitions.len()
+            );
+        }
+        Arc::make_mut(&mut self.cache).partitioning = output_partitioning;
+        Ok(self)
+    }
+
     pub fn partitions(&self) -> &Vec<Arc<dyn PartitionStream>> {
         &self.partitions
     }
@@ -147,14 +167,12 @@ impl StreamingTableExec {
     fn compute_properties(
         schema: SchemaRef,
         orderings: Vec<LexOrdering>,
-        partitions: &[Arc<dyn PartitionStream>],
+        output_partitioning: Partitioning,
         infinite: bool,
     ) -> PlanProperties {
         // Calculate equivalence properties:
         let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings);
 
-        // Get output partitioning:
-        let output_partitioning = Partitioning::UnknownPartitioning(partitions.len());
         let boundedness = if infinite {
             Boundedness::Unbounded {
                 requires_infinite_memory: false,
@@ -204,6 +222,16 @@ impl DisplayAs for StreamingTableExec {
                 if let Some(fetch) = self.limit {
                     write!(f, ", fetch={fetch}")?;
                 }
+                if !matches!(
+                    self.cache.output_partitioning(),
+                    Partitioning::UnknownPartitioning(_)
+                ) {
+                    write!(
+                        f,
+                        ", output_partitioning={}",
+                        self.cache.output_partitioning()
+                    )?;
+                }
 
                 display_orderings(f, &self.projected_output_ordering)?;
 
@@ -306,6 +334,17 @@ impl ExecutionPlan for StreamingTableExec {
             };
             lex_orderings.push(ordering);
         }
+        let projection_mapping = ProjectionMapping::try_new(
+            projection
+                .expr()
+                .iter()
+                .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())),
+            &self.schema(),
+        )?;
+        let output_partitioning = self
+            .cache
+            .output_partitioning()
+            .project(&projection_mapping, self.cache.equivalence_properties());
 
         StreamingTableExec::try_new(
             Arc::clone(self.partition_schema()),
@@ -315,6 +354,7 @@ impl ExecutionPlan for StreamingTableExec {
             self.is_infinite(),
             self.limit(),
         )
+        .and_then(|exec| exec.with_output_partitioning(output_partitioning))
         .map(|e| Some(Arc::new(e) as _))
     }