Add config-matrix tests in enforce_distribution.rs for range-satisfaction settings

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -70,6 +70,7 @@ use datafusion_physical_plan::joins::utils::JoinOn;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_physical_plan::test::exec::KeyPartitioningRequirementExec;
 use datafusion_physical_plan::union::UnionExec;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, displayable,
@@ -747,6 +748,160 @@ impl TestConfig {
     }
 }
 
+#[derive(Debug, Clone, Copy)]
+enum RangeKeyMatch {
+    Exact,
+    Subset,
+    Incompatible,
+}
+
+#[derive(Debug, Clone, Copy)]
+enum ThresholdState {
+    Met,
+    NotMet,
+}
+
+impl ThresholdState {
+    fn value(self, input_partitions: usize) -> usize {
+        match self {
+            Self::Met => input_partitions,
+            Self::NotMet => input_partitions + 1,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+enum TargetPartitions {
+    Equal,
+    Greater,
+}
+
+impl TargetPartitions {
+    fn value(self, input_partitions: usize) -> usize {
+        match self {
+            Self::Equal => input_partitions,
+            Self::Greater => input_partitions + 1,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+enum ExpectedPlan {
+    Reuse,
+    Hash,
+}
+
+#[derive(Debug, Clone, Copy)]
+struct RangeSatisfactionConfigCase {
+    subset_threshold: ThresholdState,
+    preserve_threshold: ThresholdState,
+    target_partitions: TargetPartitions,
+    // Expected plans for Exact, Subset, and Incompatible keys, respectively.
+    expected_plans: [ExpectedPlan; 3],
+}
+
+#[test]
+fn range_satisfaction_config_matrix() -> Result<()> {
+    const INPUT_PARTITIONS: usize = 4;
+    use ExpectedPlan::{Hash, Reuse};
+    use RangeKeyMatch::{Exact, Incompatible, Subset};
+    use TargetPartitions::{Equal, Greater};
+    use ThresholdState::{Met, NotMet};
+
+    let config_cases = [
+        // subset  preserve  target     exact  subset  incompatible
+        RangeSatisfactionConfigCase {
+            subset_threshold: NotMet,
+            preserve_threshold: NotMet,
+            target_partitions: Equal,
+            expected_plans: [Reuse, Hash, Hash],
+        },
+        RangeSatisfactionConfigCase {
+            subset_threshold: NotMet,
+            preserve_threshold: NotMet,
+            target_partitions: Greater,
+            expected_plans: [Hash, Hash, Hash],
+        },
+        RangeSatisfactionConfigCase {
+            subset_threshold: NotMet,
+            preserve_threshold: Met,
+            target_partitions: Equal,
+            expected_plans: [Reuse, Hash, Hash],
+        },
+        RangeSatisfactionConfigCase {
+            subset_threshold: NotMet,
+            preserve_threshold: Met,
+            target_partitions: Greater,
+            expected_plans: [Reuse, Reuse, Hash],
+        },
+        RangeSatisfactionConfigCase {
+            subset_threshold: Met,
+            preserve_threshold: NotMet,
+            target_partitions: Equal,
+            expected_plans: [Reuse, Reuse, Hash],
+        },
+        RangeSatisfactionConfigCase {
+            subset_threshold: Met,
+            preserve_threshold: NotMet,
+            target_partitions: Greater,
+            expected_plans: [Reuse, Reuse, Hash],
+        },
+        RangeSatisfactionConfigCase {
+            subset_threshold: Met,
+            preserve_threshold: Met,
+            target_partitions: Equal,
+            expected_plans: [Reuse, Reuse, Hash],
+        },
+        RangeSatisfactionConfigCase {
+            subset_threshold: Met,
+            preserve_threshold: Met,
+            target_partitions: Greater,
+            expected_plans: [Reuse, Reuse, Hash],
+        },
+    ];
+    for config_case in config_cases {
+        for (key_match, expected_plan) in [Exact, Subset, Incompatible]
+            .into_iter()
+            .zip(config_case.expected_plans)
+        {
+            let input = parquet_exec_with_output_partitioning(range_partitioning(
+                "a",
+                [10, 20, 30],
+                SortOptions::default(),
+            )?);
+            let partition_keys = match key_match {
+                Exact => vec![col("a", &schema())?],
+                Subset => vec![col("a", &schema())?, col("b", &schema())?],
+                Incompatible => vec![col("b", &schema())?],
+            };
+            let requirement =
+                Arc::new(KeyPartitioningRequirementExec::new(input, partition_keys));
+
+            let mut config = TestConfig::default().with_query_execution_partitions(
+                config_case.target_partitions.value(INPUT_PARTITIONS),
+            );
+            config.config.optimizer.subset_repartition_threshold =
+                config_case.subset_threshold.value(INPUT_PARTITIONS);
+            config.config.optimizer.preserve_file_partitions =
+                config_case.preserve_threshold.value(INPUT_PARTITIONS);
+
+            let plan = config.to_plan(requirement, &DISTRIB_DISTRIB_SORT);
+            let plan = displayable(plan.as_ref()).indent(true).to_string();
+            let has_hash_repartition =
+                plan.contains("RepartitionExec: partitioning=Hash");
+
+            assert_eq!(
+                has_hash_repartition,
+                matches!(expected_plan, Hash),
+                "unexpected optimized plan for key_match={key_match:?}, \
+                 config={config_case:?}:\n{plan}"
+            );
+        }
+    }
+
+    Ok(())
+}
+
 #[test]
 fn range_aggregate_reuses_range_partitioning() -> Result<()> {
     let input = parquet_exec_with_output_partitioning(range_partitioning(

datafusion/physical-plan/src/test/exec.rs

@@ -18,9 +18,9 @@
 //! Simple iterator over batches for use in testing
 
 use crate::{
-    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
-    RecordBatchStream, SendableRecordBatchStream, Statistics, common,
-    execution_plan::Boundedness, statistics::StatisticsArgs,
+    DisplayAs, DisplayFormatType, ExecutionPlan, InputDistributionRequirements,
+    Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream,
+    Statistics, common, execution_plan::Boundedness, statistics::StatisticsArgs,
 };
 use crate::{
     execution_plan::EmissionType,
@@ -37,7 +37,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::{DataFusionError, Result, internal_err};
 use datafusion_execution::TaskContext;
-use datafusion_physical_expr::EquivalenceProperties;
+use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalExpr};
 
 use futures::Stream;
 use tokio::sync::Barrier;
@@ -276,6 +276,85 @@ impl ExecutionPlan for MockExec {
     }
 }
 
+/// A test [`ExecutionPlan`] that requires key partitioning from its input and
+/// opts into compatible Range partitioning satisfying that requirement.
+#[derive(Debug)]
+pub struct KeyPartitioningRequirementExec {
+    input: Arc<dyn ExecutionPlan>,
+    partition_keys: Vec<Arc<dyn PhysicalExpr>>,
+    cache: Arc<PlanProperties>,
+}
+
+impl KeyPartitioningRequirementExec {
+    /// Create a new key-partitioning requirement for `input`.
+    pub fn new(
+        input: Arc<dyn ExecutionPlan>,
+        partition_keys: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> Self {
+        let cache = Arc::clone(input.properties());
+        Self {
+            input,
+            partition_keys,
+            cache,
+        }
+    }
+}
+
+impl DisplayAs for KeyPartitioningRequirementExec {
+    fn fmt_as(
+        &self,
+        t: DisplayFormatType,
+        f: &mut std::fmt::Formatter,
+    ) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "KeyPartitioningRequirementExec")
+            }
+            DisplayFormatType::TreeRender => write!(f, ""),
+        }
+    }
+}
+
+impl ExecutionPlan for KeyPartitioningRequirementExec {
+    fn name(&self) -> &'static str {
+        "KeyPartitioningRequirementExec"
+    }
+
+    fn properties(&self) -> &Arc<PlanProperties> {
+        &self.cache
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        vec![&self.input]
+    }
+
+    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
+        InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
+            self.partition_keys.clone(),
+        )])
+        .allow_range_satisfaction_for_key_partitioning()
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        mut children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        assert_eq!(children.len(), 1);
+        Ok(Arc::new(Self::new(
+            children.swap_remove(0),
+            self.partition_keys.clone(),
+        )))
+    }
+
+    fn execute(
+        &self,
+        _partition: usize,
+        _context: Arc<TaskContext>,
+    ) -> Result<SendableRecordBatchStream> {
+        unreachable!("test exec does not support execution")
+    }
+}
+
 fn clone_error(e: &DataFusionError) -> DataFusionError {
     use DataFusionError::*;
     match e {