ShayanGho commented on code in PR #25589:
URL: https://github.com/apache/datafusion/pull/25589#discussion_r4244563250
##########
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 detailed repro. Fixed in the last commit: the guard now
clears per-partition constants on a copy of the input’s equivalence properties
before checking a multi-partition input, the same way `RepartitionExec` does.
`range_join_key_alignment_ignores_per_partition_constants` covers your case: an
unbounded right input ordered `(b, a)` with a per-partition constant a keeps
its keys and passes `SanityCheckPlan`.
--
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]
