imback82 commented on a change in pull request #28676:
URL: https://github.com/apache/spark/pull/28676#discussion_r456580570
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
##########
@@ -2665,6 +2665,16 @@ object SQLConf {
.checkValue(_ > 0, "The difference must be positive.")
.createWithDefault(4)
+ val BROADCAST_HASH_JOIN_OUTPUT_PARTITIONING_EXPAND_LIMIT =
+
buildConf("spark.sql.execution.broadcastHashJoin.outputPartitioningExpandLimit")
Review comment:
Updated
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
##########
@@ -60,6 +63,81 @@ case class BroadcastHashJoinExec(
}
}
+ override lazy val outputPartitioning: Partitioning = {
+ joinType match {
+ case Inner =>
+ streamedPlan.outputPartitioning match {
+ case h: HashPartitioning => expandOutputPartitioning(h)
+ case c: PartitioningCollection => expandOutputPartitioning(c)
+ case other => other
+ }
+ case _ => streamedPlan.outputPartitioning
+ }
+ }
+
+ // An one-to-many mapping from a streamed key to build keys.
+ private lazy val streamedKeyToBuildKeyMapping = {
+ val mapping = mutable.Map.empty[Expression, Seq[Expression]]
+ streamedKeys.zip(buildKeys).foreach {
+ case (streamedKey, buildKey) =>
+ val key = streamedKey.canonicalized
+ mapping.get(key) match {
+ case Some(v) => mapping.put(key, v :+ buildKey)
+ case None => mapping.put(key, Seq(buildKey))
+ }
+ }
+ mapping.toMap
+ }
+
+ // Expands the given partitioning collection recursively.
+ private def expandOutputPartitioning(
+ partitioning: PartitioningCollection): PartitioningCollection = {
+ PartitioningCollection(partitioning.partitionings.flatMap {
+ case h: HashPartitioning => expandOutputPartitioning(h).partitionings
+ case c: PartitioningCollection => Seq(expandOutputPartitioning(c))
+ case other => Seq(other)
+ })
+ }
+
+ // Expands the given hash partitioning by substituting streamed keys with
build keys.
+ // For example, if the expressions for the given partitioning are Seq("a",
"b", "c")
+ // where the streamed keys are Seq("b", "c") and the build keys are Seq("x",
"y"),
+ // the expanded partitioning will have the following expressions:
+ // Seq("a", "b", "c"), Seq("a", "b", "y"), Seq("a", "x", "c"), Seq("a", "x",
"y").
+ // The expanded expressions are returned as PartitioningCollection.
+ private def expandOutputPartitioning(partitioning: HashPartitioning):
PartitioningCollection = {
+ val maxNumCombinations = sqlContext.conf.getConf(
+ SQLConf.BROADCAST_HASH_JOIN_OUTPUT_PARTITIONING_EXPAND_LIMIT)
+ var currentNumCombinations = 0
+
+ def generateExprCombinations(
+ current: Seq[Expression],
+ accumulated: Seq[Expression]): Seq[Seq[Expression]] = {
+ if (currentNumCombinations > maxNumCombinations) {
+ Nil
+ } else if (current.isEmpty) {
+ currentNumCombinations += 1
+ Seq(accumulated)
+ } else {
+ val buildKeys =
streamedKeyToBuildKeyMapping.get(current.head.canonicalized)
+ generateExprCombinations(current.tail, accumulated :+ current.head) ++
+ buildKeys.map { bKeys =>
+ bKeys.flatMap { bKey =>
+ if (currentNumCombinations < maxNumCombinations) {
Review comment:
Wanted to avoid unnecessary recursion (+ not creating new `Seq`, etc.),
but I removed the check for simplicity.
----------------------------------------------------------------
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]
