sunchao commented on code in PR #25589:
URL: https://github.com/apache/datafusion/pull/25589#discussion_r4209329677
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -554,6 +554,123 @@ fn reorder_current_join_keys(
}
}
+/// Finds the permutation that aligns this input's join keys with its range
keys.
+fn range_join_key_positions(
+ input: &dyn ExecutionPlan,
+ join_keys: &[PhysicalExprRef],
+) -> Option<Vec<usize>> {
+ let Partitioning::Range(range) = input.output_partitioning() else {
+ return None;
+ };
+
+ if join_keys.len() != range.ordering().len() {
+ return None;
+ }
+
+ let expected = range
+ .ordering()
+ .iter()
+ .map(|sort_expr| Arc::clone(&sort_expr.expr))
+ .collect::<Vec<_>>();
+
+ expected_expr_positions(join_keys, &expected).or_else(|| {
+ let eq_group = input.equivalence_properties().eq_group();
+ if eq_group.is_empty() {
+ return None;
+ }
+
+ let normalized_keys = join_keys
+ .iter()
+ .map(|expr| eq_group.normalize_expr(Arc::clone(expr)))
+ .collect::<Vec<_>>();
+ let normalized_expected = expected
+ .iter()
+ .map(|expr| eq_group.normalize_expr(Arc::clone(expr)))
+ .collect::<Vec<_>>();
+
+ expected_expr_positions(&normalized_keys, &normalized_expected)
+ })
+}
+
+/// Returns the non-identity permutation that aligns `on` with a
+/// range-partitioned input, trying the left input first and then the right.
+fn range_aligned_join_key_positions(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ on: &[(PhysicalExprRef, PhysicalExprRef)],
+) -> Option<Vec<usize>> {
+ if !matches!(left.output_partitioning(), Partitioning::Range(_))
+ && !matches!(right.output_partitioning(), Partitioning::Range(_))
+ {
+ return None;
+ }
+
+ let keys = extract_join_keys(on);
+ let positions = range_join_key_positions(left.as_ref(), &keys.left_keys)
+ .or_else(|| range_join_key_positions(right.as_ref(),
&keys.right_keys))?;
+
+ if positions.iter().copied().eq(0..positions.len()) {
+ return None;
+ }
+
+ Some(positions)
+}
+
+/// Aligns a partitioned join's equi-key pairs with an existing input range
+/// ordering, so compatible Range inputs satisfy the join's co-partitioning
+/// requirement without a repartition. Covers partitioned hash joins and
+/// sort-merge joins. The permutation moves whole key pairs and, for
+/// sort-merge joins, the matching `sort_options`.
+fn reorder_join_keys_to_range_inputs(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+ if let Some(join) = plan.downcast_ref::<HashJoinExec>() {
+ if join.mode != PartitionMode::Partitioned {
+ return Ok(plan);
+ }
+ let Some(positions) =
+ range_aligned_join_key_positions(&join.left, &join.right, &join.on)
+ else {
+ return Ok(plan);
+ };
+ let new_on = positions
+ .into_iter()
+ .map(|index| join.on[index].clone())
+ .collect();
+ return join.builder().with_on(new_on).build_exec();
+ }
+
+ if let Some(join) = plan.downcast_ref::<SortMergeJoinExec>() {
+ let Some(positions) =
+ range_aligned_join_key_positions(&join.left, &join.right, &join.on)
Review Comment:
Thanks for the fix. I reran the original reproducer on `3c53a04`, and that
case now passes. There is still a **[P2] streaming regression when the ordering
check relies on a per-partition constant**.
In `forces_unbounded_sort` ([lines
609–611](https://github.com/apache/datafusion/blob/3c53a047ef4c6b0e3fa97102f97e4b33f0d7c830/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs#L609-L611)),
the check uses the child's properties before distribution enforcement. If `a`
is a `ConstExpr` with `AcrossPartitions::Heterogeneous`, an input explicitly
ordered by `(b, a)` satisfies the proposed `(a, b)` locally. However, the
subsequent multi-input `RepartitionExec` calls
`clear_per_partition_constants()`. Its output properties no longer prove the
rewritten ordering, so sorting enforcement inserts a full `SortExec` on the
unbounded input.
I reproduced this with a bounded left input using two `Range(a, b)`
partitions split at `(10, 0)`, and an unbounded right input using
`UnknownPartitioning(2)`. Both advertise ordering `(b, a)` and the valid
per-partition constant `a`. Partition `p = 0, 1` contains 16 rows `a = 2 +
10*p, b = 2*i + p` for `i = 0..15`; the right streams emit their batch and then
remain pending. With target partitions set to 2 and round-robin/file-scan
repartitioning disabled:
- On merge base `6047791`, sanity validation passes and the join emits a
first batch.
- On head `3c53a04`, both `top_down_join_key_reordering` settings insert the
blocking sort: `SanityCheckPlan` rejects the plan with `Cannot execute pipeline
breaking queries`, and direct execution times out waiting one second for its
first batch.
- The same rows with the left advertised as `UnknownPartitioning(2)` stream
successfully on both revisions. Finite controls return all 32 exact expected
rows on both revisions. The six-case diagnostic passes 6/6 on base and 4/6 on
head.
Could the guard use ordering guarantees that survive the planned exchanges?
A conservative option is to clear per-partition constants on a cloned
equivalence-properties value before checking a multi-partition input. Please
include a regression covering this case before resolving the streaming thread.
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -554,6 +554,155 @@ fn reorder_current_join_keys(
}
}
+/// Finds the permutation that aligns this input's join keys with its range
keys.
+fn range_join_key_positions(
+ input: &dyn ExecutionPlan,
+ join_keys: &[PhysicalExprRef],
+) -> Option<Vec<usize>> {
+ let Partitioning::Range(range) = input.output_partitioning() else {
+ return None;
+ };
+
+ if join_keys.len() != range.ordering().len() {
Review Comment:
Confirmed on `3c53a04`: this has a concrete cost for sort-merge joins. With
`target_partitions = 1`, two bounded inputs using `Range([a, b], [])`, and both
inputs already ordered by `(b, a)`, the original `(b, a)` join is valid and
needs no exchange or sort. Alignment changes the keys to `(a, b)` and adds
**two full `SortExec` operators**, while still adding no exchange. I reproduced
this under both `top_down_join_key_reordering` settings.
Please skip alignment when the inputs will remain single-partition and there
is no distribution benefit. This is plan evidence of unnecessary sorting; I did
not measure runtime. A regression with `target_partitions = 1` would make the
intended behavior explicit.
--
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]
