cloud-fan commented on a change in pull request #29074:
URL: https://github.com/apache/spark/pull/29074#discussion_r500395857
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
##########
@@ -171,19 +176,54 @@ case class EnsureRequirements(conf: SQLConf) extends
Rule[SparkPlan] {
leftPartitioning: Partitioning,
rightPartitioning: Partitioning): (Seq[Expression], Seq[Expression]) = {
if (leftKeys.forall(_.deterministic) && rightKeys.forall(_.deterministic))
{
- (leftPartitioning, rightPartitioning) match {
- case (HashPartitioning(leftExpressions, _), _) =>
- reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq,
leftExpressions, leftKeys)
- case (_, HashPartitioning(rightExpressions, _)) =>
- reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq,
rightExpressions, rightKeys)
- case _ =>
- (leftKeys, rightKeys)
- }
+ reorderJoinKeysRecursively(
+ leftKeys,
+ rightKeys,
+ Some(leftPartitioning),
+ Some(rightPartitioning))
+ .getOrElse((leftKeys, rightKeys))
} else {
(leftKeys, rightKeys)
}
}
+ /**
+ * Recursively reorders the join keys based on partitioning. It starts
reordering the
+ * join keys to match HashPartitioning on either side, followed by
PartitioningCollection.
+ */
+ private def reorderJoinKeysRecursively(
+ leftKeys: Seq[Expression],
+ rightKeys: Seq[Expression],
+ leftPartitioning: Option[Partitioning],
+ rightPartitioning: Option[Partitioning]): Option[(Seq[Expression],
Seq[Expression])] = {
+ (leftPartitioning, rightPartitioning) match {
+ case (Some(HashPartitioning(leftExpressions, _)), _) =>
+ reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq,
leftExpressions, leftKeys)
+ .orElse(reorderJoinKeysRecursively(
+ leftKeys, rightKeys, None, rightPartitioning))
+ case (_, Some(HashPartitioning(rightExpressions, _))) =>
+ reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq,
rightExpressions, rightKeys)
+ .orElse(reorderJoinKeysRecursively(
+ leftKeys, rightKeys, leftPartitioning, None))
+ case (Some(PartitioningCollection(partitionings)), _) =>
+ partitionings.foreach { p =>
+ reorderJoinKeysRecursively(leftKeys, rightKeys, Some(p),
rightPartitioning).map { k =>
+ return Some(k)
Review comment:
nit:
```
partitionings.foldLeft(None) { (res, p) =>
res.orElse(reorderJoinKeysRecursively...)
}.getOrElse(reorderJoinKeysRecursively(leftKeys, rightKeys, None,
rightPartitioning))
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
