sunchao commented on a change in pull request #32875:
URL: https://github.com/apache/spark/pull/32875#discussion_r748558824
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
##########
@@ -352,3 +380,142 @@ case class BroadcastPartitioning(mode: BroadcastMode)
extends Partitioning {
case _ => false
}
}
+
+/**
+ * This is used in the scenario where an operator has multiple children (e.g.,
join) and one or more
+ * of which have their own requirement regarding whether its data can be
considered as
+ * co-partitioned from others. This offers APIs for:
+ *
+ * 1. Comparing with specs from other children of the operator and check if
they are compatible.
+ * When two specs are compatible, we can say their data are
co-partitioned, and Spark will
+ * potentially able to eliminate shuffle if necessary.
+ * 1. Creating a partitioning that can be used to re-partition another
child, so that to make it
+ * having a compatible partitioning as this node.
+ */
+trait ShuffleSpec {
+ /**
+ * Returns the number of partitions of this shuffle spec
+ */
+ def numPartitions: Int
+
+ /**
+ * Returns true iff this spec is compatible with the other [[Partitioning]]
and
+ * clustering expressions (e.g., from [[ClusteredDistribution]]).
+ *
+ * A true return value means that the data partitioning from this spec can
be seen as
+ * co-partitioned with the `otherPartitioning`, and therefore no shuffle is
required when
+ * joining the two sides.
+ */
+ def isCompatibleWith(other: ShuffleSpec): Boolean
+
+ /**
+ * Creates a partitioning that can be used to re-partitioned the other side
with the given
+ * required distribution.
+ *
+ * Note: this will only be called after `isCompatibleWith` returns true on
the side where the
+ * `clustering` is returned from.
+ */
+ final def createPartitioning(distribution: Distribution): Partitioning =
distribution match {
+ case AllTuples =>
+ SinglePartition
+ case ClusteredDistribution(clustering, _) =>
+ createPartitioning0(clustering)
+ case _ =>
+ throw new IllegalStateException("unexpected distribution: " +
+ s"${distribution.getClass.getSimpleName}")
+ }
+
+ def createPartitioning0(clustering: Seq[Expression]): Partitioning
+}
+
+case object SinglePartitionShuffleSpec extends ShuffleSpec {
+ override def isCompatibleWith(other: ShuffleSpec): Boolean = {
Review comment:
> If a binary plan requires its left child to be AllTuples and the
right child to be ClusteredDistribution, I don't think it makes sense to
require the children to be co-partitioned.
I think this is the case of `FlatMapCoGroupsInPandasExec` where one side
could be `AllTuples` and the other side could be `ClusteredDistribution`, in
which case we should still avoid shuffle if the latter only has one partition?
do you mean we should handle it separately, like in `EnsureRequirements` but
not in `ShuffleSpec`?
> createPartitioning should take numPartitions and keyIndexes (not an
arbitrary Distribution) and create partitioning w.r.t. its own Distribution.
This is sort of what's done at the moment. `createPartitioning` only takes
`AllTuples` and `ClusteredDistribution` and it's an error on the implementer
side if other distribution types are supplied. Subclasses of `ShuffleSpec` are
only required to implement `createPartitioning0` which takes a sequence of
expressions (e.g., cluster keys) from the side to be re-shuffled. The
numPartitions and keyIndexes are maintained by the shuffle spec (with best
parallelism) that is used to re-shuffle the other side.
For the case:
```sql
SELECT * FROM A JOIN B ON A.c1 = B.d1 AND a.c2 = B.d2
```
If `A` is `HashPartitioning(c1, 10)` and `B` is `HashPartitioning(d2, 5)`,
the shuffle spec of `A` will maintain key indexes `{ 0 }`. When re-shuffling
`B`, it'll pick `c2` from the input expressions `[c2, d2]` using the indexes
and create a partitioning `HashPartitioning(c2, 10)`.
--
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]
