refactor: one rule for the equivalence properties of a merge of partitions

datafusion/physical-expr/src/equivalence/properties/mod.rs

@@ -362,6 +362,25 @@ impl EquivalenceProperties {
         }
     }
 
+    /// The properties of the output of an operator that combines the rows of
+    /// `input_partitions` input partitions, such as `CoalescePartitionsExec`,
+    /// `SortPreservingMergeExec` and `RepartitionExec`, given the properties
+    /// of its input.
+    ///
+    /// With one input partition no rows from different partitions meet, so
+    /// everything that held still holds. Otherwise constants that may differ
+    /// between partitions are removed, and so are the orderings unless
+    /// `keeps_order` is set.
+    pub fn merged(mut self, input_partitions: usize, keeps_order: bool) -> Self {
+        if input_partitions > 1 {
+            if !keeps_order {
+                self.clear_orderings();
+            }
+            self.clear_per_partition_constants();
+        }
+        self
+    }
+
     /// Adds new orderings into the existing ordering equivalence class.
     pub fn add_orderings(
         &mut self,
@@ -1795,4 +1814,42 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_merged_keeps_per_partition_facts_of_one_partition() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int32, true),
+            Field::new("b", DataType::Int32, true),
+        ]));
+        let a = col("a", &schema)?;
+        let b = col("b", &schema)?;
+        // Each partition is sorted by `a` and has one value of `b`, which may
+        // differ between partitions.
+        let mut input = EquivalenceProperties::new(Arc::clone(&schema));
+        input.add_ordering([PhysicalSortExpr::new_default(Arc::clone(&a))]);
+        input.add_constants([ConstExpr::new(
+            Arc::clone(&b),
+            AcrossPartitions::Heterogeneous,
+        )])?;
+        let b_is_constant = |p: &EquivalenceProperties| p.is_expr_constant(&b).is_some();
+        let sorted_by_a = |p: &EquivalenceProperties| p.output_ordering().is_some();
+
+        // One input partition: no rows from different partitions meet.
+        for keeps_order in [false, true] {
+            let merged = input.clone().merged(1, keeps_order);
+            assert!(b_is_constant(&merged), "keeps_order={keeps_order}");
+            assert!(sorted_by_a(&merged), "keeps_order={keeps_order}");
+        }
+
+        // Several input partitions: `b` may take several values.
+        let merged = input.clone().merged(2, true);
+        assert!(!b_is_constant(&merged));
+        assert!(sorted_by_a(&merged));
+
+        let merged = input.merged(2, false);
+        assert!(!b_is_constant(&merged));
+        assert!(!sorted_by_a(&merged));
+
+        Ok(())
+    }
 }

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -606,12 +606,12 @@ fn forces_unbounded_sort<'a>(
     let sort_exprs = keys
         .zip(sort_options)
         .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr), *options));
-    // A repartition of a multi-partition input drops its per-partition
-    // constants, so do not rely on them to satisfy the ordering.
-    let mut eq_properties = input.equivalence_properties().clone();
-    if input.output_partitioning().partition_count() > 1 {
-        eq_properties.clear_per_partition_constants();
-    }
+    // An order-preserving repartition of a multi-partition input drops its
+    // per-partition constants, so do not rely on them to satisfy the ordering.
+    let eq_properties = input
+        .equivalence_properties()
+        .clone()
+        .merged(input.output_partitioning().partition_count(), true);
     Ok(!eq_properties.ordering_satisfy(sort_exprs)?)
 }
 

datafusion/physical-plan/src/coalesce_partitions.rs

@@ -93,10 +93,11 @@ impl CoalescePartitionsExec {
             )
         };
 
-        // Coalescing partitions loses existing orderings:
-        let mut eq_properties = input.equivalence_properties().clone();
-        eq_properties.clear_orderings();
-        eq_properties.clear_per_partition_constants();
+        // Coalescing several partitions loses their orderings:
+        let eq_properties = input
+            .equivalence_properties()
+            .clone()
+            .merged(input_partitions, false);
         PlanProperties::new(
             eq_properties,                        // Equivalence Properties
             Partitioning::UnknownPartitioning(1), // Output Partitioning
@@ -469,13 +470,15 @@ mod tests {
     use crate::test::exec::{
         BarrierExec, BlockingExec, PanicExec, assert_strong_count_converges_to_zero,
     };
-    use crate::test::{self, assert_is_pending};
+    use crate::test::{self, TestMemoryExec, assert_is_pending};
     use crate::{collect, common};
 
     use std::time::Duration;
 
     use arrow::array::RecordBatch;
     use arrow::datatypes::{DataType, Field, Schema};
+    use datafusion_physical_expr::expressions::col;
+    use datafusion_physical_expr_common::sort_expr::LexOrdering;
 
     use futures::FutureExt;
 
@@ -705,4 +708,27 @@ mod tests {
             CardinalityEffect::LowerEqual
         ));
     }
+
+    #[test]
+    fn keeps_input_ordering_of_one_partition() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
+        let ordering: LexOrdering =
+            [PhysicalSortExpr::new_default(col("a", &schema)?)].into();
+        let input = |partitions: usize| -> Result<Arc<dyn ExecutionPlan>> {
+            let batches = vec![vec![]; partitions];
+            let source = TestMemoryExec::try_new(&batches, Arc::clone(&schema), None)?
+                .try_with_sort_information(vec![ordering.clone()])?;
+            Ok(Arc::new(TestMemoryExec::update_cache(&Arc::new(source))))
+        };
+
+        // One input partition is passed through, so its order still holds.
+        let coalesce = CoalescePartitionsExec::new(input(1)?);
+        assert_eq!(coalesce.properties().output_ordering(), Some(&ordering));
+
+        // Rows of several partitions are interleaved.
+        let coalesce = CoalescePartitionsExec::new(input(2)?);
+        assert_eq!(coalesce.properties().output_ordering(), None);
+
+        Ok(())
+    }
 }

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

@@ -2200,18 +2200,10 @@ impl RepartitionExec {
         input: &Arc<dyn ExecutionPlan>,
         preserve_order: bool,
     ) -> EquivalenceProperties {
-        // Equivalence Properties
-        let mut eq_properties = input.equivalence_properties().clone();
-        // If the ordering is lost, reset the ordering equivalence class:
-        if !Self::maintains_input_order_helper(input, preserve_order)[0] {
-            eq_properties.clear_orderings();
-        }
-        // When there are more than one input partitions, they will be fused at the output.
-        // Therefore, remove per partition constants.
-        if input.output_partitioning().partition_count() > 1 {
-            eq_properties.clear_per_partition_constants();
-        }
-        eq_properties
+        input.equivalence_properties().clone().merged(
+            input.output_partitioning().partition_count(),
+            preserve_order,
+        )
     }
 
     /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.

datafusion/physical-plan/src/sorts/sort_preserving_merge.rs

@@ -171,8 +171,10 @@ impl SortPreservingMergeExec {
             )
         };
 
-        let mut eq_properties = input.equivalence_properties().clone();
-        eq_properties.clear_per_partition_constants();
+        let mut eq_properties = input
+            .equivalence_properties()
+            .clone()
+            .merged(input_partitions, true);
         eq_properties.add_ordering(ordering);
         PlanProperties::new(
             eq_properties,                        // Equivalence Properties