imback82 commented on a change in pull request #28676:
URL: https://github.com/apache/spark/pull/28676#discussion_r455505839
##########
File path:
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
##########
@@ -60,6 +62,67 @@ case class BroadcastHashJoinExec(
}
}
+ override lazy val outputPartitioning: Partitioning = {
+ joinType match {
+ case _: InnerLike =>
+ 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 = {
+ def generateExprCombinations(
+ current: Seq[Expression],
+ accumulated: Seq[Expression]): Seq[Seq[Expression]] = {
+ if (current.isEmpty) {
+ Seq(accumulated)
+ } else {
+ val buildKeys =
streamedKeyToBuildKeyMapping.get(current.head.canonicalized)
+ generateExprCombinations(current.tail, accumulated :+ current.head) ++
+ buildKeys.map { _.flatMap(b =>
generateExprCombinations(current.tail, accumulated :+ b))
Review comment:
I added a config to limit the expansion. (Please let me know if
introducing a new config knob is too much. Then I can just have a constant in
this class - not configurable).
----------------------------------------------------------------
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:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
