diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index c280a0aa5ef80..974b9da617119 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -385,6 +385,22 @@ fn parquet_exec_multiple_sorted( fn parquet_exec_with_output_partitioning( output_partitioning: Partitioning, ) -> Arc { + parquet_exec_with_partitioning_and_size( + output_partitioning, + Precision::Absent, + Precision::Absent, + ) +} + +fn parquet_exec_with_partitioning_and_size( + output_partitioning: Partitioning, + num_rows: Precision, + total_byte_size: Precision, +) -> Arc { + let mut statistics = Statistics::new_unknown(&schema()); + statistics.num_rows = num_rows; + statistics.total_byte_size = total_byte_size; + let file_groups = (0..output_partitioning.partition_count()) .map(|partition| { FileGroup::new(vec![PartitionedFile::new(format!("p{partition}"), 100)]) @@ -397,6 +413,7 @@ fn parquet_exec_with_output_partitioning( ) .with_file_groups(file_groups) .with_output_partitioning(Some(output_partitioning)) + .with_statistics(statistics) .build(); DataSourceExec::from_data_source(config) @@ -1123,6 +1140,42 @@ fn range_hash_join_repartitions_unpartitioned_side_to_match_range() -> Result<() Ok(()) } +#[test] +fn range_hash_join_rejects_smaller_reference() -> Result<()> { + let left = parquet_exec_with_partitioning_and_size( + Partitioning::UnknownPartitioning(4), + Precision::Inexact(100_000), + Precision::Inexact(8_000_000), + ); + let right = parquet_exec_with_partitioning_and_size( + range_partitioning("a", [10, 20, 30], SortOptions::default())?, + Precision::Inexact(100), + Precision::Inexact(8_000), + ); + + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + let children = plan.children(); + assert_eq!(children.len(), 2); + for child in children { + assert!( + matches!(child.output_partitioning(), Partitioning::Hash(_, 4)), + "expected hash fallback, got {:?}", + child.output_partitioning(), + ); + } + + Ok(()) +} + #[test] fn range_requirement_scales_within_sample_resolution() -> Result<()> { let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into(); @@ -1170,6 +1223,76 @@ fn range_requirement_scales_within_sample_resolution() -> Result<()> { Ok(()) } +#[test] +fn range_hash_join_preserves_equal_reference() -> Result<()> { + let left = parquet_exec_with_partitioning_and_size( + Partitioning::UnknownPartitioning(4), + Precision::Inexact(100), + Precision::Inexact(8_000), + ); + let right = parquet_exec_with_partitioning_and_size( + range_partitioning("a", [10, 20, 30], SortOptions::default())?, + Precision::Inexact(100), + Precision::Inexact(8_000), + ); + + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + let expected = range_partitioning("a", [10, 20, 30], SortOptions::default())?; + + let children = plan.children(); + assert_eq!(children.len(), 2); + for child in children { + assert_eq!(child.output_partitioning(), &expected); + } + + Ok(()) +} + +#[test] +fn range_hash_join_preserves_large_reference() -> Result<()> { + // Both inputs are above the `CollectLeft` thresholds, so the slightly + // larger one is adapted to the reference and only it is shuffled. + let left = parquet_exec_with_partitioning_and_size( + Partitioning::UnknownPartitioning(4), + Precision::Inexact(11_000_000), + Precision::Inexact(880_000_000), + ); + let right = parquet_exec_with_partitioning_and_size( + range_partitioning("a", [10, 20, 30], SortOptions::default())?, + Precision::Inexact(10_000_000), + Precision::Inexact(800_000_000), + ); + + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + let expected = range_partitioning("a", [10, 20, 30], SortOptions::default())?; + + let children = plan.children(); + assert_eq!(children.len(), 2); + for child in children { + assert_eq!(child.output_partitioning(), &expected); + } + + Ok(()) +} + #[test] fn range_singleton_scaling_respects_required_keys() -> Result<()> { let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into(); diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 3c2426a5ef2b1..ead9463355c1e 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -1078,6 +1078,21 @@ impl PlanSize { fn is_known(&self) -> bool { self.byte_size.is_some() || self.num_rows.is_some() } + + /// Whether the plan is below the `CollectLeft` thresholds of `JoinSelection`, + /// comparing bytes when known and rows otherwise. `None` when the size is unknown. + fn is_small(&self, config: &ConfigOptions) -> Option { + let optimizer = &config.optimizer; + match (self.byte_size, self.num_rows) { + (Some(bytes), _) => { + Some(bytes < optimizer.hash_join_single_partition_threshold) + } + (None, Some(rows)) => { + Some(rows < optimizer.hash_join_single_partition_threshold_rows) + } + (None, None) => None, + } + } } impl PartialOrd for PlanSize { @@ -1123,8 +1138,9 @@ fn enforce_distribution_relationships( plan_name: &str, input_distributions: &InputDistributionRequirements, children: &mut [DistributionChildState], - target_partitions: usize, + config: &ConfigOptions, ) -> Result<()> { + let target_partitions = config.execution.target_partitions; let mut repartitioned_for_relationship = vec![false; children.len()]; loop { @@ -1210,6 +1226,24 @@ fn enforce_distribution_relationships( } }; + // Do not adapt a child above the `CollectLeft` thresholds to a reference + // below them; fall back to hash partitioning instead. Otherwise keep the + // reference, so that it is not shuffled as well. + let best_satisfied_child = best_satisfied_child.filter(|(ref_idx, _)| { + let reference_size = + PlanSize::from_plan(children[*ref_idx].context.plan.as_ref()); + if reference_size.is_small(config) != Some(true) { + return true; + } + + !unsatisfied_children.iter().any(|&child_idx| { + child_idx != *ref_idx + && PlanSize::from_plan(children[child_idx].context.plan.as_ref()) + .is_small(config) + == Some(false) + }) + }); + // Validate that best_satisfied_child can be adapted across all unsatisfied children in the group let best_satisfied_child = best_satisfied_child.filter(|(ref_idx, ref_partitioning)| { @@ -1665,7 +1699,7 @@ pub fn ensure_distribution_with_stats( plan.name(), &input_distributions, &mut children, - target_partitions, + config, )?; let children = children