jayzhan211 commented on code in PR #25787:
URL: https://github.com/apache/datafusion/pull/25787#discussion_r4492771963
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1362,6 +1378,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, _)| {
Review Comment:
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.
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -1976,6 +1993,42 @@ fn
range_hash_join_repartitions_unpartitioned_side_to_match_range() -> Result<()
Ok(())
}
+#[test]
+fn range_hash_join_rejects_smaller_reference() -> Result<()> {
Review Comment:
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:
```rs
#[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(())
}
```
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -1230,6 +1230,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> {
Review Comment:
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:
```diff
--- 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))
});
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
