ShayanGho commented on code in PR #25589:
URL: https://github.com/apache/datafusion/pull/25589#discussion_r4077287461
##########
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:
Fixed. SMJ alignment now proceeds only when every unbounded input already
satisfies the realigned ordering, checked through its equivalence properties.
Joins with two bounded inputs retain the existing behavior.
The regression test covers both cases under both reordering settings: inputs
ordered by (b, a) keep the original keys and use order-preserving repartitions;
inputs already ordered by (a, b) still get aligned without sorting or
repartitioning. Both pass sanity validation and remain stable across repeated
optimizer passes.
--
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]
