allow range to satisfy key distribution generally

Author: gene-bordegaray

datafusion/physical-expr/src/partitioning.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`].
-// <https://github.com/apache/datafusion/issues/23266>.
-//
-// 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<dyn ExecutionPlan>,
-    child: &Arc<dyn ExecutionPlan>,
     allow_subset_satisfy_partitioning: bool,
 ) -> bool {
     let Some(aggregate) = plan.downcast_ref::<AggregateExec>() 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,
                     );
 

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.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<dyn ExecutionPlan>) -> bool {
     plan.is::<RepartitionExec>()
 }
 
-/// TODO: remove once Range generally satisfies KeyPartitioned requirements
-/// through Partitioning::satisfaction.
-/// See <https://github.com/apache/datafusion/issues/23266>.
-///
-/// 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<dyn PhysicalExpr>],
-    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::<Vec<_>>();
-
-            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::<Vec<_>>();
-            let normalized_required_exprs = required_exprs
-                .iter()
-                .map(|expr| eq_group.normalize_expr(Arc::clone(expr)))
-                .collect::<Vec<_>>();
-
-            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<dyn ExecutionPlan>) -> bool {

datafusion/physical-optimizer/src/utils.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<Option<OrderingRequirements>> {