fix: consider input sizes before reusing join partitioning

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<DataSourceExec> {
+    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<usize>,
+    total_byte_size: Precision<usize>,
+) -> Arc<DataSourceExec> {
+    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();

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<bool> {
+        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