diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index ed824e5712f96..5fc577a91c6c9 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/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(()) + } } diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 168eef838d43a..23243cef8a421 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/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)?) } diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 7097cc780de03..46d06923f543d 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/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> { + 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(()) + } } diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index b4984ab5b234f..5580380965d28 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -2200,18 +2200,10 @@ impl RepartitionExec { input: &Arc, 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. diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index e865fef98fee7..34417f740e1ed 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/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