diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 3292ada0a8e86..2e1ec123ff224 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/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( diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index b92008c6b219b..494a95be8cd24 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/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, + partition_keys: Vec>, + cache: Arc, +} + +impl KeyPartitioningRequirementExec { + /// Create a new key-partitioning requirement for `input`. + pub fn new( + input: Arc, + partition_keys: Vec>, + ) -> 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 { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + 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, + mut children: Vec>, + ) -> Result> { + 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, + ) -> Result { + unreachable!("test exec does not support execution") + } +} + fn clone_error(e: &DataFusionError) -> DataFusionError { use DataFusionError::*; match e {