sunchao commented on code in PR #25589:
URL: https://github.com/apache/datafusion/pull/25589#discussion_r4074507963
##########
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:
**[P2] Preserve streaming order when aligning sort-merge join keys**
Could we preserve the original join-key order when this alignment would
require a blocking sort on an unbounded input? `Range(a, b)` describes
partition membership, but rows within each partition can already be sorted by
`(b, a)`. For an SMJ with keys `(b, a)`, this rewrite changes its required
ordering to `(a, b)` and introduces blocking `SortExec` operators, making a
previously runnable streaming plan fail `SanityCheckPlan` with `Cannot execute
pipeline breaking queries`.
I reproduced this through the physical execution API with two valid
`Range(a, b)` partitions split at `(10, 0)`, each emitting a batch ordered by
`(b, a)` and then remaining pending. The batch in partition `p` contains `b =
0..15` and `a = 10*p + 2 - (b % 2)`. On base `0576a0b`, the plan retains `(b,
a)` and uses order-preserving hash repartitions: sanity validation passes and
the join emits a matching row. On head `76494e1`, both children get `(a, b)`
sorts: sanity validation fails and direct execution produces no first batch
within one second. This reproduces with `top_down_join_key_reordering` both
enabled and disabled. Finite versions of the same inputs return all 32 expected
rows on both revisions.
Please retain the streamable key order in this case and add an
unbounded-input regression test.
--
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]
