Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SPARK-32282][SQL] Improve EnsureRquirement.reorderJoinKeys to handle more scenarios such as PartitioningCollection #29074

Closed
wants to merge 13 commits into from
Original file line number Diff line number Diff line change
Expand Up @@ -130,9 +130,14 @@ case class EnsureRequirements(conf: SQLConf) extends Rule[SparkPlan] {
leftKeys: IndexedSeq[Expression],
rightKeys: IndexedSeq[Expression],
expectedOrderOfKeys: Seq[Expression],
currentOrderOfKeys: Seq[Expression]): (Seq[Expression], Seq[Expression]) = {
currentOrderOfKeys: Seq[Expression]): Option[(Seq[Expression], Seq[Expression])] = {
if (expectedOrderOfKeys.size != currentOrderOfKeys.size) {
return (leftKeys, rightKeys)
return None
}

// Check if the current order already satisfies the expected order.
if (expectedOrderOfKeys.zip(currentOrderOfKeys).forall(p => p._1.semanticEquals(p._2))) {
return Some(leftKeys, rightKeys)
}

// Build a lookup between an expression and the positions its holds in the current key seq.
Expand All @@ -159,10 +164,10 @@ case class EnsureRequirements(conf: SQLConf) extends Rule[SparkPlan] {
rightKeysBuffer += rightKeys(index)
case _ =>
// The expression cannot be found, or we have exhausted all indices for that expression.
return (leftKeys, rightKeys)
return None
}
}
(leftKeysBuffer.toSeq, rightKeysBuffer.toSeq)
Some(leftKeysBuffer.toSeq, rightKeysBuffer.toSeq)
}

private def reorderJoinKeys(
Expand All @@ -171,19 +176,50 @@ 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, leftPartitioning, 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: Partitioning,
rightPartitioning: Partitioning): Option[(Seq[Expression], Seq[Expression])] = {
(leftPartitioning, rightPartitioning) match {
case (HashPartitioning(leftExpressions, _), _) =>
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leftExpressions, leftKeys)
.orElse(reorderJoinKeysRecursively(
leftKeys, rightKeys, UnknownPartitioning(0), rightPartitioning))
imback82 marked this conversation as resolved.
Show resolved Hide resolved
case (_, HashPartitioning(rightExpressions, _)) =>
reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, rightExpressions, rightKeys)
.orElse(reorderJoinKeysRecursively(
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This can be also implemented by looking at left partitioning first then move to the right partitionoing:

    (leftPartitioning, rightPartitioning) match {
      case (HashPartitioning(leftExpressions, _), _) =>
        reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leftExpressions, leftKeys)
          .orElse(reorderJoinKeysRecursively(
            leftKeys, rightKeys, UnknownPartitioning(0), rightPartitioning))
      case (PartitioningCollection(partitionings), _) =>
        partitionings.foreach { p =>
          reorderJoinKeysRecursively(leftKeys, rightKeys, p, rightPartitioning).map { k =>
            return Some(k)
          }
        }
        reorderJoinKeysRecursively(leftKeys, rightKeys, UnknownPartitioning(0), rightPartitioning)
      case (_, HashPartitioning(rightExpressions, _)) =>
        reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, rightExpressions, rightKeys)
      case (_, PartitioningCollection(partitionings)) =>
        partitionings.foreach { p =>
          reorderJoinKeysRecursively(leftKeys, rightKeys, leftPartitioning, p).map { k =>
            return Some(k)
          }
        }
        None
      case _ =>
        None
    }

However, I chose this way so that the behavior remains the same. If you have leftPartitioning = PartitioningCollection and rightPartitioning = HashPartitioning, it will match the rightPartitioning first, which is the existing behavior.

leftKeys, rightKeys, leftPartitioning, UnknownPartitioning(0)))
case (PartitioningCollection(partitionings), _) =>
partitionings.foreach { p =>
reorderJoinKeysRecursively(leftKeys, rightKeys, p, rightPartitioning).map { k =>
return Some(k)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit:

partitionings.foldLeft(None) { (res, p) =>
  res.orElse(reorderJoinKeysRecursively...)
}.getOrElse(reorderJoinKeysRecursively(leftKeys, rightKeys, None, rightPartitioning))

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, updated.

}
}
reorderJoinKeysRecursively(leftKeys, rightKeys, UnknownPartitioning(0), rightPartitioning)
case (_, PartitioningCollection(partitionings)) =>
partitionings.foreach { p =>
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you do the same refactor here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed.

reorderJoinKeysRecursively(leftKeys, rightKeys, leftPartitioning, p).map { k =>
return Some(k)
}
}
None
case _ =>
None
}
}

/**
* When the physical operators are created for JOIN, the ordering of join keys is based on order
* in which the join keys appear in the user query. That might not match with the output
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.spark.sql.execution.exchange

import org.apache.spark.sql.catalyst.expressions.Literal
import org.apache.spark.sql.catalyst.plans.Inner
import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, PartitioningCollection}
import org.apache.spark.sql.execution.{DummySparkPlan, SortExec}
import org.apache.spark.sql.execution.joins.SortMergeJoinExec
import org.apache.spark.sql.test.SharedSparkSession

class EnsureRequirementsSuite extends SharedSparkSession {
private val exprA = Literal(1)
private val exprB = Literal(2)
private val exprC = Literal(3)

test("EnsureRequirements.reorder should handle PartitioningCollection") {
imback82 marked this conversation as resolved.
Show resolved Hide resolved
val plan1 = DummySparkPlan(
outputPartitioning = PartitioningCollection(Seq(
HashPartitioning(exprA :: exprB :: Nil, 5),
HashPartitioning(exprA :: Nil, 5))))
val plan2 = DummySparkPlan()

// Test PartitioningCollection on the left side of join.
val smjExec1 = SortMergeJoinExec(
exprB :: exprA :: Nil, exprA :: exprB :: Nil, Inner, None, plan1, plan2)
EnsureRequirements(spark.sessionState.conf).apply(smjExec1) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _,
DummySparkPlan(_, _, PartitioningCollection(leftPartitionings), _, _), _),
SortExec(_, _,
ShuffleExchangeExec(HashPartitioning(rightPartitioningExpressions, _), _, _), _), _) =>
assert(leftKeys !== smjExec1.leftKeys)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this check needed? We already check leftKeys === Seq(exprA, exprB) and it's obvious that leftKeys !== smjExec1.leftKeys

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed, thanks!

assert(rightKeys !== smjExec1.rightKeys)
assert(leftKeys === leftPartitionings.head.asInstanceOf[HashPartitioning].expressions)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we simply check leftKeys === Seq(exprA, exprB)?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK. I simplified the checks in this test.

assert(rightKeys === rightPartitioningExpressions)
case other => fail(other.toString)
}

// Test PartitioningCollection on the right side of join.
val smjExec2 = SortMergeJoinExec(
exprA :: exprB :: Nil, exprB :: exprA :: Nil, Inner, None, plan2, plan1)
EnsureRequirements(spark.sessionState.conf).apply(smjExec2) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _,
ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _), _),
SortExec(_, _,
DummySparkPlan(_, _, PartitioningCollection(rightPartitionings), _, _), _), _) =>
assert(leftKeys !== smjExec2.leftKeys)
assert(rightKeys !== smjExec2.rightKeys)
assert(leftKeys === leftPartitioningExpressions)
assert(rightKeys === rightPartitionings.head.asInstanceOf[HashPartitioning].expressions)
case other => fail(other.toString)
}

// Both sides are PartitioningCollection, but left side cannot be reorderd to match
// and it should fall back to the right side.
val smjExec3 = SortMergeJoinExec(
exprA :: exprC :: Nil, exprB :: exprA :: Nil, Inner, None, plan1, plan1)
EnsureRequirements(spark.sessionState.conf).apply(smjExec3) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _,
ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _), _),
SortExec(_, _,
DummySparkPlan(_, _, PartitioningCollection(rightPartitionings), _, _), _), _) =>
assert(leftKeys !== smjExec3.leftKeys)
assert(rightKeys !== smjExec3.rightKeys)
assert(leftKeys === leftPartitioningExpressions)
assert(rightKeys === rightPartitionings.head.asInstanceOf[HashPartitioning].expressions)
case other => fail(other.toString)
}
}

test("EnsureRequirements.reorder should fallback to the other side partitioning") {
imback82 marked this conversation as resolved.
Show resolved Hide resolved
val plan1 = DummySparkPlan(
outputPartitioning = HashPartitioning(exprA :: exprB :: exprC :: Nil, 5))
val plan2 = DummySparkPlan(
outputPartitioning = HashPartitioning(exprB :: exprC :: Nil, 5))

// Test fallback to the right side, which has PartitioningCollection.
imback82 marked this conversation as resolved.
Show resolved Hide resolved
val smjExec1 = SortMergeJoinExec(
exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan2)
EnsureRequirements(spark.sessionState.conf).apply(smjExec1) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _,
ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _), _),
SortExec(_, _,
DummySparkPlan(_, _, HashPartitioning(rightPartitioningExpressions, _), _, _), _), _) =>
assert(leftKeys !== smjExec1.leftKeys)
assert(rightKeys !== smjExec1.rightKeys)
assert(leftKeys === leftPartitioningExpressions)
assert(rightKeys === rightPartitioningExpressions)
case other => fail(other.toString)
}

// Test fallback to the right side, which has PartitioningCollection.
val plan3 = DummySparkPlan(
outputPartitioning = PartitioningCollection(Seq(HashPartitioning(exprB :: exprC :: Nil, 5))))
val smjExec2 = SortMergeJoinExec(
exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan3)
EnsureRequirements(spark.sessionState.conf).apply(smjExec2) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _,
ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _), _),
SortExec(_, _,
DummySparkPlan(_, _, PartitioningCollection(rightPartitionings), _, _), _), _) =>
assert(leftKeys !== smjExec2.leftKeys)
assert(rightKeys !== smjExec2.rightKeys)
assert(leftKeys === leftPartitioningExpressions)
assert(rightKeys === rightPartitionings.head.asInstanceOf[HashPartitioning].expressions)
case other => fail(other.toString)
}

// The right side has HashPartitioning, so it is matched first, but no reordering match is
// found, and it should fall back to the left side, which has a PartitioningCollection.
val smjExec3 = SortMergeJoinExec(
exprC :: exprB :: Nil, exprA :: exprB :: Nil, Inner, None, plan3, plan1)
EnsureRequirements(spark.sessionState.conf).apply(smjExec3) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _,
DummySparkPlan(_, _, PartitioningCollection(leftPartitionings), _, _), _),
SortExec(_, _,
ShuffleExchangeExec(HashPartitioning(rightPartitioningExpressions, _), _, _), _), _) =>
assert(leftKeys !== smjExec3.leftKeys)
assert(rightKeys !== smjExec3.rightKeys)
assert(leftKeys === leftPartitionings.head.asInstanceOf[HashPartitioning].expressions)
assert(rightKeys === rightPartitioningExpressions)
case other => fail(other.toString)
}
}
}