ShayanGho commented on code in PR #25589:
URL: https://github.com/apache/datafusion/pull/25589#discussion_r4077112454
##########
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 @sunchao, good catch! I missed the unbounded-input case when adding
SMJ alignment. I’ll reproduce this and guard the rewrite so it preserves the
existing key order when alignment would introduce a blocking sort. I’ll also
add a regression test under both reordering settings.
--
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]
