From de0572d5071e14ffc8fb0e9aedf3d73d0f482d1a Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Fri, 17 Jul 2026 16:02:40 -0400 Subject: [PATCH] allow range to satisfy key distribution generally --- datafusion/physical-expr/src/partitioning.rs | 433 +++++++----------- .../enforce_distribution.rs | 37 +- datafusion/physical-optimizer/src/utils.rs | 59 +-- .../physical-plan/src/aggregates/mod.rs | 12 +- .../src/distribution_requirements.rs | 165 +------ .../physical-plan/src/joins/hash_join/exec.rs | 8 +- .../src/joins/sort_merge_join/exec.rs | 1 - .../src/joins/symmetric_hash_join.rs | 1 - .../src/sorts/partitioned_topk.rs | 1 - .../src/windows/bounded_window_agg_exec.rs | 1 - .../src/windows/window_agg_exec.rs | 1 - 11 files changed, 182 insertions(+), 537 deletions(-) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 61492934ebd19..85c5c95330c84 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -417,40 +417,26 @@ impl Partitioning { // Here we do not check the partition count for hash partitioning and assumes the partition count // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins, // then we need to have the partition count and hash functions validation. - Partitioning::Hash(partition_exprs, _) => { - // Empty hash partitioning is invalid - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { - return PartitioningSatisfaction::Exact; - } - - let eq_groups = eq_properties.eq_group(); - if !eq_groups.is_empty() { - if allow_subset { - let normalized_partition_exprs = - normalize_exprs(partition_exprs, eq_properties); - let normalized_required_exprs = - normalize_exprs(required_exprs, eq_properties); - if Self::is_subset_partitioning( - &normalized_partition_exprs, - &normalized_required_exprs, - ) { - return PartitioningSatisfaction::Subset; - } - } - } else if allow_subset - && Self::is_subset_partitioning(partition_exprs, required_exprs) - { - return PartitioningSatisfaction::Subset; - } - - PartitioningSatisfaction::NotSatisfied + Partitioning::Hash(partition_exprs, _) => Self::key_satisfaction( + partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ), + Partitioning::Range(range) => { + let partition_exprs = range + .ordering() + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + Self::key_satisfaction( + &partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ) } Partitioning::RoundRobinBatch(_) - | Partitioning::Range(_) | Partitioning::UnknownPartitioning(_) => { PartitioningSatisfaction::NotSatisfied } @@ -459,6 +445,43 @@ impl Partitioning { } } + fn key_satisfaction( + partition_exprs: &[Arc], + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, + ) -> PartitioningSatisfaction { + if partition_exprs.is_empty() || required_exprs.is_empty() { + return PartitioningSatisfaction::NotSatisfied; + } + + if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { + return PartitioningSatisfaction::Exact; + } + + let eq_groups = eq_properties.eq_group(); + if !eq_groups.is_empty() { + if allow_subset { + let normalized_partition_exprs = + normalize_exprs(partition_exprs, eq_properties); + let normalized_required_exprs = + normalize_exprs(required_exprs, eq_properties); + if Self::is_subset_partitioning( + &normalized_partition_exprs, + &normalized_required_exprs, + ) { + return PartitioningSatisfaction::Subset; + } + } + } else if allow_subset + && Self::is_subset_partitioning(partition_exprs, required_exprs) + { + return PartitioningSatisfaction::Subset; + } + + PartitioningSatisfaction::NotSatisfied + } + /// Calculate the output partitioning after applying the given projection. pub fn project( &self, @@ -685,6 +708,26 @@ mod tests { } } + fn assert_satisfaction( + desc: &str, + partitioning: &Partitioning, + required: &Distribution, + eq_properties: &EquivalenceProperties, + expected_with_subset: PartitioningSatisfaction, + expected_without_subset: PartitioningSatisfaction, + ) { + assert_eq!( + partitioning.satisfaction(required, eq_properties, true), + expected_with_subset, + "Failed for {desc} with subset enabled" + ); + assert_eq!( + partitioning.satisfaction(required, eq_properties, false), + expected_without_subset, + "Failed for {desc} with subset disabled" + ); + } + #[test] #[expect( deprecated, @@ -768,320 +811,114 @@ mod tests { } #[test] - fn test_partitioning_satisfy_by_subset() -> Result<()> { + fn test_partitioning_key_satisfaction() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); let test_cases = vec![ ( - "KeyPartitioned([a, b]) satisfied by Hash([a])", - fixture.hash_partitioning([0], 4), + "exact: KeyPartitioned([a, b]) satisfied by Hash([a, b])", + fixture.hash_partitioning([0, 1], 4), fixture.key_distribution([0, 1]), - PartitioningSatisfaction::Subset, - PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, ), ( - "KeyPartitioned([a, b, c]) satisfied by Hash([a])", + "subset: KeyPartitioned([a, b]) satisfied by Hash([a])", fixture.hash_partitioning([0], 4), - fixture.key_distribution([0, 1, 2]), - PartitioningSatisfaction::Subset, - PartitioningSatisfaction::NotSatisfied, - ), - ( - "KeyPartitioned([a, b, c]) satisfied by Hash([a, b])", - fixture.hash_partitioning([0, 1], 4), - fixture.key_distribution([0, 1, 2]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a, b, c]) satisfied by Hash([b])", + "subset: KeyPartitioned([a, b, c]) satisfied by Hash([b])", fixture.hash_partitioning([1], 4), fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", + "subset reordered: KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", fixture.hash_partitioning([1, 0], 4), fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &fixture.eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_current_superset() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; - - let test_cases = vec![ ( - "KeyPartitioned([a]) satisfied by Hash([a, b])", + "superset: KeyPartitioned([a]) not satisfied by Hash([a, b])", fixture.hash_partitioning([0, 1], 4), fixture.key_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a]) satisfied by Hash([a, b, c])", + "superset: KeyPartitioned([a, b]) not satisfied by Hash([a, b, c])", fixture.hash_partitioning([0, 1, 2], 4), - fixture.key_distribution([0]), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([a, b]) satisfied by Hash([a, b, c])", - fixture.hash_partitioning([0, 1, 2], 4), + "partial overlap: KeyPartitioned([a, b]) not satisfied by Hash([a, c])", + fixture.hash_partitioning([0, 2], 4), fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &fixture.eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_partial_overlap() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; - - let test_cases = vec![( - "Partial overlap: KeyPartitioned([a, b]) satisfied by Hash([a, c])", - fixture.hash_partitioning([0, 2], 4), - fixture.key_distribution([0, 1]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - )]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &fixture.eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_no_overlap() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; - - let test_cases = vec![ ( - "KeyPartitioned([b, c]) satisfied by Hash([a])", + "no overlap: KeyPartitioned([b, c]) not satisfied by Hash([a])", fixture.hash_partitioning([0], 4), fixture.key_distribution([1, 2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([c]) satisfied by Hash([a, b])", - fixture.hash_partitioning([0, 1], 4), - fixture.key_distribution([2]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &fixture.eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_exact_match() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - - let test_cases = vec![ - ( - "KeyPartitioned([a, b]) satisfied by Hash([a, b])", - fixture.hash_partitioning([0, 1], 4), - fixture.key_distribution([0, 1]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, - ), - ( - "KeyPartitioned([a]) satisfied by Hash([a])", - fixture.hash_partitioning([0], 4), - fixture.key_distribution([0]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, - ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &fixture.eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_unknown() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); - - let test_cases = vec![ - ( - "KeyPartitioned([a, b]) satisfied by Hash([unknown])", + "unknown partition expr", Partitioning::Hash(vec![Arc::clone(&unknown)], 4), fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([unknown]) satisfied by Hash([a, b])", + "unknown required expr", fixture.hash_partitioning([0, 1], 4), Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([unknown]) satisfied by Hash([unknown])", + "same unknown expr", Partitioning::Hash(vec![Arc::clone(&unknown)], 4), Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([unknown, a]) satisfied by Hash([unknown])", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::KeyPartitioned(vec![Arc::clone(&unknown), fixture.col(0)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &fixture.eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_empty_hash() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a"])?; - - let test_cases = vec![ - ( - "KeyPartitioned([a]) satisfied by Hash([])", + "empty hash partitioning", Partitioning::Hash(vec![], 4), fixture.key_distribution([0]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "KeyPartitioned([]) satisfied by Hash([a])", + "empty key distribution", fixture.hash_partitioning([0], 4), Distribution::KeyPartitioned(vec![]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), - ( - "KeyPartitioned([]) satisfied by Hash([])", - Partitioning::Hash(vec![], 4), - Distribution::KeyPartitioned(vec![]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), ]; for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &fixture.eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &fixture.eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" + assert_satisfaction( + desc, + &partition, + &required, + &fixture.eq_properties, + expected_with_subset, + expected_without_subset, ); } @@ -1230,15 +1067,65 @@ mod tests { } #[test] - fn test_multi_partition_range_does_not_satisfy_hash_distribution() -> Result<()> { - let fixture = PartitioningTestFixture::int64(&["a", "b"])?; - let range_partitioning = + fn test_range_partitioning_key_satisfaction() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let range_a = fixture.range_partitioning([0], vec![int_split_point([10])]); + let range_ab = fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); - let required = fixture.key_distribution([0, 1]); - assert_eq!( - range_partitioning.satisfaction(&required, &fixture.eq_properties, false), - PartitioningSatisfaction::NotSatisfied + assert_satisfaction( + "exact single key", + &range_a, + &fixture.key_distribution([0]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "exact compound key", + &range_ab, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "subset key", + &range_a, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + assert_satisfaction( + "incompatible key", + &range_a, + &fixture.key_distribution([1]), + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(2))?; + assert_satisfaction( + "equivalent subset key", + &range_a, + &fixture.key_distribution([1, 2]), + &eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?; + assert_satisfaction( + "equivalent exact key", + &range_a, + &fixture.key_distribution([1]), + &eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, ); Ok(()) diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 55769fce01c06..952aae9846d0f 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -35,7 +35,7 @@ use std::sync::Arc; use crate::output_requirements::OutputRequirementExec; use crate::utils::{ add_sort_above_with_check, is_coalesce_partitions, is_repartition, - is_sort_preserving_merge, range_partitioning_satisfies_key_partitioning, + is_sort_preserving_merge, }; use arrow::compute::SortOptions; @@ -698,18 +698,13 @@ fn add_roundrobin_on_top( } } -// TODO: remove this temporary bridge once [`Partitioning::Range`] -// generally satisfies [`Distribution::KeyPartitioned`] through -// [`Partitioning::satisfaction`]. -// . -// -// Partial aggregates do not require key partitioning, but they preserve their -// input partitioning for the final aggregate. Until Range satisfies -// KeyPartitioned generally, this check keeps preserve_file_partitions from -// inserting RoundRobin between a reusable Range input and the partial aggregate. -fn partial_aggregate_preserves_reusable_partitioning( +// Partial aggregates require unspecified input distribution, but their output +// may already satisfy the final aggregate's key distribution because partial +// aggregation preserves/projects input partitioning. Keep that reusable output +// partitioning intact when preserve_file_partitions would otherwise insert +// RoundRobin below the partial aggregate. +fn partial_aggregate_output_satisfies_final_partitioning( plan: &Arc, - child: &Arc, allow_subset_satisfy_partitioning: bool, ) -> bool { let Some(aggregate) = plan.downcast_ref::() else { @@ -722,24 +717,15 @@ fn partial_aggregate_preserves_reusable_partitioning( return false; } - let group_exprs = aggregate.group_expr().input_exprs(); - let output_partitioning = child.output_partitioning(); - let eq_properties = child.equivalence_properties(); - let key_distribution = Distribution::KeyPartitioned(group_exprs.clone()); + let key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr()); - output_partitioning + plan.output_partitioning() .satisfaction( &key_distribution, - eq_properties, + plan.equivalence_properties(), allow_subset_satisfy_partitioning, ) .is_satisfied() - || range_partitioning_satisfies_key_partitioning( - output_partitioning, - &group_exprs, - eq_properties, - allow_subset_satisfy_partitioning, - ) } /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator @@ -1308,9 +1294,8 @@ pub fn ensure_distribution( let preserve_partial_aggregate_partitioning = preserve_file_partition_threshold_met - && partial_aggregate_preserves_reusable_partitioning( + && partial_aggregate_output_satisfies_final_partitioning( &plan, - &child.plan, allow_subset_satisfy_partitioning, ); diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 1fbf8c6fb78cd..04229e1cc2737 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -18,10 +18,7 @@ use std::sync::Arc; use datafusion_common::Result; -use datafusion_physical_expr::{ - Distribution, EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, - PhysicalExpr, physical_exprs_equal, -}; +use datafusion_physical_expr::{Distribution, LexOrdering, LexRequirement}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -161,60 +158,6 @@ pub fn is_repartition(plan: &Arc) -> bool { plan.is::() } -/// TODO: remove once Range generally satisfies KeyPartitioned requirements -/// through Partitioning::satisfaction. -/// See . -/// -/// Checks whether range partitioning satisfies a key partitioning requirement. -/// This is intentionally separate from general partitioning satisfaction while -/// range reuse is rolled out operator by operator. -pub(crate) fn range_partitioning_satisfies_key_partitioning( - partitioning: &Partitioning, - required_exprs: &[Arc], - eq_properties: &EquivalenceProperties, - allow_subset: bool, -) -> bool { - match partitioning { - Partitioning::Range(range) => { - let partition_exprs = range - .ordering() - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - - if partition_exprs.is_empty() || required_exprs.is_empty() { - return false; - } - - let eq_group = eq_properties.eq_group(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - let normalized_required_exprs = required_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - - if physical_exprs_equal( - &normalized_required_exprs, - &normalized_partition_exprs, - ) { - return true; - } - - allow_subset - && normalized_partition_exprs.len() < normalized_required_exprs.len() - && normalized_partition_exprs.iter().all(|partition_expr| { - normalized_required_exprs - .iter() - .any(|required_expr| partition_expr.eq(required_expr)) - }) - } - _ => false, - } -} - /// Checks whether the given operator is a limit; /// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`]. pub fn is_limit(plan: &Arc) -> bool { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 732da32ab0391..30e0ad24695b2 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1934,7 +1934,7 @@ impl ExecutionPlan for AggregateExec { } fn input_distribution_requirements(&self) -> InputDistributionRequirements { - let requirements = InputDistributionRequirements::new(match &self.mode { + InputDistributionRequirements::new(match &self.mode { AggregateMode::Partial | AggregateMode::PartialReduce => { vec![Distribution::UnspecifiedDistribution] } @@ -1944,15 +1944,7 @@ impl ExecutionPlan for AggregateExec { AggregateMode::Final | AggregateMode::Single => { vec![Distribution::SinglePartition] } - }); - match &self.mode { - AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned - if !self.group_by.has_grouping_set() => - { - requirements.allow_range_satisfaction_for_key_partitioning() - } - _ => requirements, - } + }) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs index 9c7a1336c06a3..6405b1f121ef7 100644 --- a/datafusion/physical-plan/src/distribution_requirements.rs +++ b/datafusion/physical-plan/src/distribution_requirements.rs @@ -17,13 +17,8 @@ //! Input distribution requirements for physical execution plans. -use std::sync::Arc; - use datafusion_common::{Result, internal_err}; -use datafusion_physical_expr::{ - Distribution, EquivalenceProperties, Partitioning, PartitioningSatisfaction, - PhysicalExpr, physical_exprs_equal, -}; +use datafusion_physical_expr::{Distribution, Partitioning, PartitioningSatisfaction}; use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; @@ -98,10 +93,7 @@ impl InputDistributionRequirements { pub fn new(per_child: Vec) -> Self { let children = per_child .into_iter() - .map(|distribution| ChildDistributionRequirement { - distribution, - satisfaction: InputDistributionSatisfaction::Default, - }) + .map(|distribution| ChildDistributionRequirement { distribution }) .collect(); Self { @@ -176,8 +168,7 @@ impl InputDistributionRequirements { ); }; - Ok(requirement.satisfaction.satisfaction( - child.output_partitioning(), + Ok(child.output_partitioning().satisfaction( &requirement.distribution, child.equivalence_properties(), options.allow_subset(), @@ -208,30 +199,6 @@ impl InputDistributionRequirements { Ok(co_partitioned.clone()) } - /// TODO: remove this temporary bridge once [`Partitioning::Range`] - /// generally satisfies [`Distribution::KeyPartitioned`] through - /// [`Partitioning::satisfaction`]. - /// . - /// - /// Also allow compatible [`Partitioning::Range`] to satisfy - /// [`Distribution::KeyPartitioned`]. - #[expect( - deprecated, - reason = "HashPartitioned is accepted during the KeyPartitioned migration" - )] - pub(crate) fn allow_range_satisfaction_for_key_partitioning(mut self) -> Self { - for child in &mut self.children { - if matches!( - child.distribution, - Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) - ) { - child.satisfaction = - InputDistributionSatisfaction::AllowRangeKeyPartitioning; - } - } - self - } - /// Validate the requirements against a plan's children. pub(crate) fn check_invariants( &self, @@ -301,10 +268,8 @@ impl InputDistributionRequirements { let first = children[first_idx]; let first_partitioning = first.output_partitioning(); - if !first_requirement - .satisfaction + if !first_partitioning .satisfaction( - first_partitioning, &first_requirement.distribution, first.equivalence_properties(), false, @@ -317,19 +282,16 @@ impl InputDistributionRequirements { for &child_idx in co_partitioned.iter().skip(1) { let requirement = &self.children[child_idx]; let child = children[child_idx]; - if !requirement - .satisfaction + if !child + .output_partitioning() .satisfaction( - child.output_partitioning(), &requirement.distribution, child.equivalence_properties(), false, ) .is_satisfied() || !compatible_co_partitioning_layout( - first_requirement, first_partitioning, - requirement, child.output_partitioning(), ) { @@ -345,60 +307,6 @@ impl InputDistributionRequirements { #[derive(Debug, Clone)] struct ChildDistributionRequirement { distribution: Distribution, - satisfaction: InputDistributionSatisfaction, -} - -/// TODO: remove this temporary bridge once [`Partitioning::Range`] -/// generally satisfies [`Distribution::KeyPartitioned`] through -/// [`Partitioning::satisfaction`]. -/// . -#[non_exhaustive] -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -enum InputDistributionSatisfaction { - /// Use [`Partitioning::satisfaction`] as-is. - #[default] - Default, - /// Also allow [`Partitioning::Range`] to satisfy - /// [`Distribution::KeyPartitioned`]. - AllowRangeKeyPartitioning, -} - -impl InputDistributionSatisfaction { - /// Returns how `partitioning` satisfies `requirement`. - #[expect( - deprecated, - reason = "HashPartitioned is accepted during the KeyPartitioned migration" - )] - fn satisfaction( - self, - partitioning: &Partitioning, - requirement: &Distribution, - eq_properties: &EquivalenceProperties, - allow_subset: bool, - ) -> PartitioningSatisfaction { - let satisfaction = - partitioning.satisfaction(requirement, eq_properties, allow_subset); - if satisfaction.is_satisfied() { - return satisfaction; - } - - if !matches!(self, Self::AllowRangeKeyPartitioning) { - return PartitioningSatisfaction::NotSatisfied; - } - - let (Distribution::HashPartitioned(required_exprs) - | Distribution::KeyPartitioned(required_exprs)) = requirement - else { - return PartitioningSatisfaction::NotSatisfied; - }; - - range_satisfies_key_partitioning( - partitioning, - required_exprs, - eq_properties, - allow_subset, - ) - } } fn validate_child_index( @@ -421,62 +329,8 @@ fn validate_child_index( Ok(()) } -/// TODO: remove this temporary bridge once [`Partitioning::Range`] -/// generally satisfies [`Distribution::KeyPartitioned`] through -/// [`Partitioning::satisfaction`]. -/// . -fn range_satisfies_key_partitioning( - partitioning: &Partitioning, - required_exprs: &[Arc], - eq_properties: &EquivalenceProperties, - allow_subset: bool, -) -> PartitioningSatisfaction { - let Partitioning::Range(range) = partitioning else { - return PartitioningSatisfaction::NotSatisfied; - }; - - let partition_exprs = range - .ordering() - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - let eq_group = eq_properties.eq_group(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - let normalized_required_exprs = required_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - - if physical_exprs_equal(&normalized_required_exprs, &normalized_partition_exprs) { - return PartitioningSatisfaction::Exact; - } - - if allow_subset - && normalized_partition_exprs.len() < normalized_required_exprs.len() - && normalized_partition_exprs.iter().all(|partition_expr| { - normalized_required_exprs - .iter() - .any(|required_expr| partition_expr.eq(required_expr)) - }) - { - PartitioningSatisfaction::Subset - } else { - PartitioningSatisfaction::NotSatisfied - } -} - fn compatible_co_partitioning_layout( - first: &ChildDistributionRequirement, first_partitioning: &Partitioning, - other: &ChildDistributionRequirement, other_partitioning: &Partitioning, ) -> bool { if first_partitioning.partition_count() == 1 @@ -491,12 +345,7 @@ fn compatible_co_partitioning_layout( match (first_partitioning, other_partitioning) { (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, - (Partitioning::Range(left), Partitioning::Range(right)) - if first.satisfaction - == InputDistributionSatisfaction::AllowRangeKeyPartitioning - && other.satisfaction - == InputDistributionSatisfaction::AllowRangeKeyPartitioning => - { + (Partitioning::Range(left), Partitioning::Range(right)) => { left.split_points() == right.split_points() && left.ordering().len() == right.ordering().len() && left diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index c622663f03a56..c746aba028990 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1271,7 +1271,7 @@ impl ExecutionPlan for HashJoinExec { } fn input_distribution_requirements(&self) -> InputDistributionRequirements { - let requirements = match self.mode { + match self.mode { PartitionMode::Partitioned => { let (left_expr, right_expr) = self .on @@ -1291,12 +1291,6 @@ impl ExecutionPlan for HashJoinExec { Distribution::UnspecifiedDistribution, Distribution::UnspecifiedDistribution, ]), - }; - - if self.mode == PartitionMode::Partitioned && self.join_type == JoinType::Inner { - requirements.allow_range_satisfaction_for_key_partitioning() - } else { - requirements } } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 5d3621c49219b..83af24528b047 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -424,7 +424,6 @@ impl ExecutionPlan for SortMergeJoinExec { Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), ]) - .allow_range_satisfaction_for_key_partitioning() } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index eb781e633c638..b95e5e1e3d493 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -427,7 +427,6 @@ impl ExecutionPlan for SymmetricHashJoinExec { Distribution::KeyPartitioned(left_expr), Distribution::KeyPartitioned(right_expr), ]) - .allow_range_satisfaction_for_key_partitioning() } StreamJoinPartitionMode::SinglePartition => { InputDistributionRequirements::new(vec![ diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index c250130341dc1..730440a429c68 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -355,7 +355,6 @@ impl ExecutionPlan for PartitionedTopKExec { crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( partition_exprs, )]) - .allow_range_satisfaction_for_key_partitioning() } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 97cafd24c0a0d..3ca612bbdb775 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -333,7 +333,6 @@ impl ExecutionPlan for BoundedWindowAggExec { InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( self.partition_keys(), )]) - .allow_range_satisfaction_for_key_partitioning() } } diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 4e8dbc06f09a9..3eb8edd298901 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -241,7 +241,6 @@ impl ExecutionPlan for WindowAggExec { InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( self.partition_keys(), )]) - .allow_range_satisfaction_for_key_partitioning() } }