fix: consider input sizes before reusing join partitioning

datafusion/common/src/config.rs

@@ -1849,11 +1849,15 @@ config_namespace! {
         pub enable_piecewise_merge_join: bool, default = false
 
         /// The maximum estimated size in bytes for one input side of a HashJoin
-        /// will be collected into a single partition
+        /// will be collected into a single partition. Partitioned joins also use
+        /// it: an input below it is not used as the co-partitioning reference for
+        /// an input above it.
         pub hash_join_single_partition_threshold: usize, default = 4 * 1024 * 1024
 
         /// The maximum estimated size in rows for one input side of a HashJoin
-        /// will be collected into a single partition
+        /// will be collected into a single partition. Partitioned joins also use
+        /// it: an input below it is not used as the co-partitioning reference for
+        /// an input above it.
         pub hash_join_single_partition_threshold_rows: usize, default = 1024 * 128
 
         /// Maximum size in bytes for the build side of a hash join to be pushed down as an InList expression for dynamic filtering.

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -380,6 +380,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)])
@@ -392,6 +408,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)
@@ -1976,6 +1993,116 @@ fn range_hash_join_repartitions_unpartitioned_side_to_match_range() -> Result<()
     Ok(())
 }
 
+#[test]
+fn range_join_reference_size_gate() -> Result<()> {
+    let optimizer = TestConfig::default().config.optimizer;
+    let bytes = optimizer.hash_join_single_partition_threshold;
+    let rows = optimizer.hash_join_single_partition_threshold_rows;
+    let known = |n: usize| Precision::Inexact(n);
+    let absent = Precision::Absent;
+
+    // (name, other (rows, bytes), reference (rows, bytes), expect hash fallback)
+    let cases = [
+        (
+            "bytes: straddle",
+            (absent, known(bytes)),
+            (absent, known(bytes - 1)),
+            true,
+        ),
+        (
+            "bytes: both at or above",
+            (absent, known(bytes + 1)),
+            (absent, known(bytes)),
+            false,
+        ),
+        (
+            "bytes: both below",
+            (absent, known(100)),
+            (absent, known(100)),
+            false,
+        ),
+        (
+            "rows: straddle",
+            (known(rows), absent),
+            (known(rows - 1), absent),
+            true,
+        ),
+        (
+            "rows: both at or above",
+            (known(rows + 1), absent),
+            (known(rows), absent),
+            false,
+        ),
+        (
+            "bytes win over rows",
+            (known(rows * 10), known(100)),
+            (known(1), known(100)),
+            false,
+        ),
+        (
+            "other unknown",
+            (absent, absent),
+            (known(1), known(100)),
+            false,
+        ),
+        (
+            "reference unknown",
+            (known(rows * 10), known(bytes * 10)),
+            (absent, absent),
+            false,
+        ),
+    ];
+
+    type MakeJoin = fn(
+        Arc<dyn ExecutionPlan>,
+        Arc<dyn ExecutionPlan>,
+        &JoinOn,
+        &JoinType,
+    ) -> Arc<dyn ExecutionPlan>;
+    let joins: [(&str, MakeJoin); 2] = [
+        ("hash", hash_join_exec),
+        ("sort-merge", sort_merge_join_exec),
+    ];
+
+    for (join_name, make_join) in joins {
+        for (name, other, reference, expect_hash) in cases {
+            let left = parquet_exec_with_partitioning_and_size(
+                Partitioning::UnknownPartitioning(4),
+                other.0,
+                other.1,
+            );
+            let right = parquet_exec_with_partitioning_and_size(
+                range_partitioning("a", [10, 20, 30], SortOptions::default())?,
+                reference.0,
+                reference.1,
+            );
+            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 = make_join(left, right, &join_on, &JoinType::Inner);
+
+            let plan = TestConfig::default()
+                .with_query_execution_partitions(4)
+                .to_plan(join, &DISTRIB_DISTRIB_SORT);
+
+            for child in plan.children() {
+                let is_hash =
+                    matches!(child.output_partitioning(), Partitioning::Hash(_, 4));
+                let is_range =
+                    matches!(child.output_partitioning(), Partitioning::Range(_));
+                assert!(
+                    if expect_hash { is_hash } else { is_range },
+                    "{join_name}, {name}: got {:?}",
+                    child.output_partitioning(),
+                );
+            }
+        }
+    }
+
+    Ok(())
+}
+
 #[test]
 fn range_requirement_scales_within_sample_resolution() -> Result<()> {
     let ordering = [PhysicalSortExpr::new_default(col("a", &schema())?)].into();

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -19,6 +19,7 @@ use std::any::Any;
 use std::fmt::Debug;
 use std::sync::Arc;
 
+use crate::join_selection::below_collect_thresholds;
 use crate::output_requirements::OutputRequirementExec;
 use crate::utils::{
     add_sort_above_with_check, is_coalesce_partitions, is_repartition,
@@ -1275,9 +1276,10 @@ fn enforce_distribution_relationships(
     plan_name: &str,
     input_distributions: &InputDistributionRequirements,
     children: &mut [DistributionChildState],
-    target_partitions: usize,
+    config: &ConfigOptions,
     stats_ctx: &StatisticsContext,
 ) -> Result<()> {
+    let target_partitions = config.execution.target_partitions;
     let mut repartitioned_for_relationship = vec![false; children.len()];
 
     loop {
@@ -1362,6 +1364,26 @@ 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 is_small = |idx: usize| {
+                let stats = stats_ctx
+                    .compute_arc(&children[idx].context.plan, &StatisticsArgs::new())
+                    .ok()?;
+                below_collect_thresholds(
+                    &stats,
+                    config.optimizer.hash_join_single_partition_threshold,
+                    config.optimizer.hash_join_single_partition_threshold_rows,
+                )
+            };
+            is_small(*ref_idx) != Some(true)
+                || !unsatisfied_children
+                    .iter()
+                    .any(|&child_idx| is_small(child_idx) == 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)| {
@@ -1841,7 +1863,7 @@ pub fn ensure_distribution_with_stats(
         plan.name(),
         &input_distributions,
         &mut children,
-        target_partitions,
+        config,
         stats_ctx,
     )?;
 

datafusion/physical-optimizer/src/join_selection.rs

@@ -111,12 +111,24 @@ fn supports_collect_by_thresholds(
     // `Exact(0)` and `Inexact(0)` are both valid stats, and we should not treat
     // them as unknown, `Absent` will return None (this is in regards to why
     // `!=0` is not checked)
+    below_collect_thresholds(&stats, threshold_byte_size, threshold_num_rows)
+        .unwrap_or(false)
+}
+
+/// Whether `stats` is below the `CollectLeft` thresholds, comparing bytes when
+/// known and rows otherwise. `None` when the size is unknown.
+pub(crate) fn below_collect_thresholds(
+    stats: &Statistics,
+    threshold_byte_size: usize,
+    threshold_num_rows: usize,
+) -> Option<bool> {
     if let Some(byte_size) = stats.total_byte_size.get_value() {
-        *byte_size < threshold_byte_size
-    } else if let Some(num_rows) = stats.num_rows.get_value() {
-        *num_rows < threshold_num_rows
+        Some(*byte_size < threshold_byte_size)
     } else {
-        false
+        stats
+            .num_rows
+            .get_value()
+            .map(|num_rows| *num_rows < threshold_num_rows)
     }
 }
 

datafusion/sqllogictest/test_files/information_schema.slt

@@ -488,8 +488,8 @@ datafusion.optimizer.expand_views_at_output false When set to true, if the retur
 datafusion.optimizer.filter_null_join_keys false When set to true, the optimizer will insert filters before a join between a nullable and non-nullable column to filter out nulls on the nullable side. This filter can add additional overhead when the file format does not fully support predicate push down.
 datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values 150 Maximum number of distinct values (rows) in the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides with more rows than this will use hash table lookups instead. Set to 0 to always use hash table lookups. This provides an additional limit beyond `hash_join_inlist_pushdown_max_size` to prevent very large IN lists that might not provide much benefit over hash table lookups. This uses the deduplicated row count once the build side has been evaluated. The default is 150 values per partition. This is inspired by Trino's `max-filter-keys-per-column` setting. See: <https://trino.io/docs/current/admin/dynamic-filtering.html#dynamic-filter-collection-thresholds>
 datafusion.optimizer.hash_join_inlist_pushdown_max_size 131072 Maximum size in bytes for the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides larger than this will use hash table lookups instead. Set to 0 to always use hash table lookups. InList pushdown can be more efficient for small build sides because it can result in better statistics pruning as well as use any bloom filters present on the scan side. InList expressions are also more transparent and easier to serialize over the network in distributed uses of DataFusion. On the other hand InList pushdown requires making a copy of the data and thus adds some overhead to the build side and uses more memory. This setting is per-partition, so we may end up using `hash_join_inlist_pushdown_max_size` * `target_partitions` memory. The default is 128kB per partition. This should allow point lookup joins (e.g. joining on a unique primary key) to use InList pushdown in most cases but avoids excessive memory usage or overhead for larger joins.
-datafusion.optimizer.hash_join_single_partition_threshold 4194304 The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition
-datafusion.optimizer.hash_join_single_partition_threshold_rows 131072 The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition
+datafusion.optimizer.hash_join_single_partition_threshold 4194304 The maximum estimated size in bytes for one input side of a HashJoin will be collected into a single partition. Partitioned joins also use it: an input below it is not used as the co-partitioning reference for an input above it.
+datafusion.optimizer.hash_join_single_partition_threshold_rows 131072 The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition. Partitioned joins also use it: an input below it is not used as the co-partitioning reference for an input above it.
 datafusion.optimizer.join_reordering true When set to true, the physical plan optimizer may swap join inputs based on statistics. When set to false, statistics-driven join input reordering is disabled and the original join order in the query is used.
 datafusion.optimizer.max_passes 3 The logical optimizer applies rules in order for up to `max_passes` passes, stopping early if a pass leaves the plan unchanged or repeats an earlier plan. The physical optimizer always runs in one pass.
 datafusion.optimizer.prefer_existing_sort false When true, DataFusion will opportunistically remove sorts when the data is already sorted, (i.e. setting `preserve_order` to true on `RepartitionExec`  and using `SortPreservingMergeExec`) When false, DataFusion will maximize plan parallelism using `RepartitionExec` even if this requires subsequently resorting data using a `SortExec`.