[ShayanGho]

Which issue does this PR close?

Part of #25302.

Rationale for this change

A small range-partitioned input can become the reference for a much larger join input, forcing the larger input to adopt the smaller input's partition boundaries.

What changes are included in this PR?

Reject a selected reference when it is below the JoinSelection CollectLeft thresholds (hash_join_single_partition_threshold{,_rows}) and another input in the co-partitioned group is above them. This allows the existing hash-partitioning fallback to apply.

Inputs on the same side of the thresholds, or with unknown sizes, keep the reference.

Auto hash joins don't change, because a small side was already CollectLeft; this fires for explicit Partitioned hash joins, sort-merge and symmetric joins.

What is the testing strategy for this PR?

Added a table test covering the byte and row thresholds, their boundaries, unknown sizes, and both hash and sort-merge joins.

Are there any user-facing changes?

Join plans can use hash repartitioning instead of adapting a larger input to a smaller input's existing partitioning.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.80%. Comparing base (6047791) to head (1aab317).
⚠️ Report is 95 commits behind head on main.

Files with missing lines	Patch %	Lines
...er/src/ensure_requirements/enforce_distribution.rs	92.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25787      +/-   ##
==========================================
+ Coverage   82.64%   82.80%   +0.15%     
==========================================
  Files        1147     1149       +2     
  Lines      445883   452257    +6374     
  Branches   445883   452257    +6374     
==========================================
+ Hits       368513   374499    +5986     
+ Misses      54982    54958      -24     
- Partials    22388    22800     +412     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[ShayanGho]

Hello @jayzhan211 @gene-bordegaray, would either of you have time for this one too? It's a 13-line change in enforce_distribution_relationships for #25302: don't adopt a reference child's partitioning when another input is strictly larger. Merged with main, and it passes alongside the new native-reference test from apache/datafusion#24766.

[jayzhan211]

Thanks @ShayanGho , 1 suggestion

[jayzhan211]

When this fires, the larger child is shuffled in both plans (it is unsatisfied), so rejecting the reference never saves a shuffle; it only adds a shuffle of the reference. With strict >, that also happens when the sizes are close: at 10M vs 11M rows, 21M rows get shuffled instead of 11M. It also reverses #24766's native-reference rule, and range_fallback_preserves_target_parallelism now asserts the opposite of what it was added to guard. Could we reject only when the reference is small, reusing JoinSelection's CollectLeft thresholds (hash_join_single_partition_threshold{,_rows})? The issue's 100-row case would still fall back to hash, inputs of similar size would keep a single shuffle, and #24766's test could keep its original assertion.

[ShayanGho]

Makes sense, done. The reference is now rejected only when it is below the CollectLeft thresholds and another input is above them. Inputs on the same side of the thresholds keep the reference, so only one side is shuffled (range_hash_join_preserves_large_reference covers the 10M vs 11M case), and range_fallback_preserves_target_parallelism is back to its original assertion.

[github-actions]

Marked as draft after 14 days without activity: merge conflicts.

Please address these items and mark Ready for review when ready.

[jayzhan211]

Thanks @ShayanGho, the plans look right on everything I probed. One test request and a few notes inline.

Could you change Closes #25302 to Part of #25302 in the description? This covers references below the thresholds, but references at or above them (4 MB by default) are still adopted, so the issue should stay open for those. It's also worth noting there that Auto hash joins don't change, because a small side was already CollectLeft; this fires for explicit Partitioned hash joins, sort-merge and symmetric joins.

In datafusion/common/src/config.rs around line 1853, these thresholds now also steer partitioned joins, so could you mention that (optional)? dev/update_config_docs.sh regenerates the docs.

-        /// The maximum estimated size in bytes for one input side of a HashJoin
-        /// will be collected into a single partition
+        /// 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.

[jayzhan211]

Could you add coverage for the rest of the gate in this PR? Removing the rows-only arm, treating an unknown input as large, or changing < to <= leaves every test passing, including the slt files, and sort-merge joins aren't tested at all. This table test fails on each of those changes and could replace the three tests here:

#[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(())
}

[jayzhan211]

Fine to handle in a follow-up under #25302, since comparing the big side's key min/max to the split points is its own PR. Size alone can't tell whether the ranges suit the big input. With a sort-merge join on 20M rows against an 8-row reference, this gate gives 183 vs 525 ms when all keys land in one range, but 156 vs 114 ms and 175 vs 145 ms when the ranges are already balanced. The slowdown comes from fewer early stops per partition plus hash skew, not from shuffling the reference.

[jayzhan211]

Optional cleanup: this repeats the bytes-then-rows rule in supports_collect_by_thresholds, and that rule has changed before (#20730). A shared helper keeps them in step and also lets the child_idx != *ref_idx check go. I ran this with the existing tests and clippy:

--- a/datafusion/physical-optimizer/src/join_selection.rs
+++ b/datafusion/physical-optimizer/src/join_selection.rs
-    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
-    } else {
-        false
-    }
-}
+    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() {
+        Some(*byte_size < threshold_byte_size)
+    } else {
+        stats
+            .num_rows
+            .get_value()
+            .map(|num_rows| *num_rows < threshold_num_rows)
+    }
+}
--- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
+++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs
         let best_satisfied_child = best_satisfied_child.filter(|(ref_idx, _)| {
-            let reference_size =
-                PlanSize::from_plan(&children[*ref_idx].context.plan, stats_ctx);
-            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, stats_ctx)
-                        .is_small(config)
-                        == Some(false)
-            })
+            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))
         });

[ShayanGho]

@jayzhan211 made the updates:

• Replaced the three tests with your table test. It fails on each of the three changes you listed.
• Applied the below_collect_thresholds helper and removed PlanSize::is_small.
• Updated both threshold descriptions (bytes and rows), regenerated configs.md, and updated information_schema.slt.

The description now says Part of #25302 and notes that Auto hash joins are unchanged. I'll take the min/max vs split-point comparison as a follow-up under #25302.