peter-toth commented on code in PR #58531:
URL: https://github.com/apache/spark/pull/58531#discussion_r3951686888
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -501,24 +521,38 @@ case class EnsureRequirements(
var newLeft = left
var newRight = right
- val specs = Seq(left, right).zip(requiredChildDistribution).map { case (p,
d) =>
- if (!d.isInstanceOf[ClusteredDistribution]) return None
- val cd = d.asInstanceOf[ClusteredDistribution]
- val specOpt = createKeyedShuffleSpec(p.outputPartitioning, cd)
- if (specOpt.isEmpty) return None
- specOpt.get
- }
-
- val leftSpec = specs.head
- val rightSpec = specs(1)
+ def candidatesFor(plan: SparkPlan, required: Distribution):
Seq[KeyedShuffleSpec] =
+ required match {
+ case cd: ClusteredDistribution =>
createKeyedShuffleSpecs(plan.outputPartitioning, cd)
+ case _ => Nil
+ }
+ val leftCandidates = candidatesFor(left, requiredChildDistribution.head)
+ val rightCandidates = candidatesFor(right, requiredChildDistribution(1))
+ if (leftCandidates.isEmpty || rightCandidates.isEmpty) return None
+
+ // Each side may offer several members, and the right one is the one the
other side can pair
+ // with, which neither side can tell on its own. So pick the pair rather
than a member per side,
+ // and rank the pairs that agree on the keys by the parallelism they
offer, the same trade
+ // `ensureDistributionAndOrdering` makes between children when it picks
`bestSpecOpt`.
+ val agreeingPairs = for {
+ l <- leftCandidates
+ r <- rightCandidates
+ if l.areKeysCompatible(r)
+ } yield (l, r)
+ val (leftSpec, rightSpec) = agreeingPairs
+ .maxByOption { case (l, r) => l.numPartitions.max(r.numPartitions) }
Review Comment:
Taken, and it went past the comment. The comment now says what the rank
measures and where it stops matching, and the arms where the merge drops a
side's keys are ranked on the side that survives.
Two things I found while chasing it.
`mergeAndDedupPartitions` has three more arms that are not the union. Under
partition filtering a `LeftOuter`, `LeftAnti`, `LeftSingle` or existence join
keeps the left's keys, and a `RightOuter` keeps the right's. There `max` reads
the side whose keys are about to be dropped, so it is wrong systematically
rather than by luck. Measured on a `LeftOuter`: the pair that wins on `max`
leaves the join on 4 groups where the old head-of-list pick happened to get 6.
`rank` now reads only the surviving side on those arms, and there is a test for
it.
The union is not a lower bound either. `reduceKeys` runs between the pick
and the merge, so a `bucket(16)` side reduced onto `bucket(8)` brings 8 keys to
a merge its spec ranked at 16.
The intersection you named stays as the one gap I am not closing, for the
reason you give. It needs the merged count, and merging every candidate pair
would also raise `storagePartitionJoinIncompatibleReducedTypesError` for pairs
that are never chosen.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -501,24 +521,38 @@ case class EnsureRequirements(
var newLeft = left
var newRight = right
- val specs = Seq(left, right).zip(requiredChildDistribution).map { case (p,
d) =>
- if (!d.isInstanceOf[ClusteredDistribution]) return None
- val cd = d.asInstanceOf[ClusteredDistribution]
- val specOpt = createKeyedShuffleSpec(p.outputPartitioning, cd)
- if (specOpt.isEmpty) return None
- specOpt.get
- }
-
- val leftSpec = specs.head
- val rightSpec = specs(1)
+ def candidatesFor(plan: SparkPlan, required: Distribution):
Seq[KeyedShuffleSpec] =
+ required match {
+ case cd: ClusteredDistribution =>
createKeyedShuffleSpecs(plan.outputPartitioning, cd)
+ case _ => Nil
+ }
+ val leftCandidates = candidatesFor(left, requiredChildDistribution.head)
+ val rightCandidates = candidatesFor(right, requiredChildDistribution(1))
+ if (leftCandidates.isEmpty || rightCandidates.isEmpty) return None
+
+ // Each side may offer several members, and the right one is the one the
other side can pair
+ // with, which neither side can tell on its own. So pick the pair rather
than a member per side,
+ // and rank the pairs that agree on the keys by the parallelism they
offer, the same trade
+ // `ensureDistributionAndOrdering` makes between children when it picks
`bestSpecOpt`.
+ val agreeingPairs = for {
Review Comment:
Right on both counts. The tie-break is in and the description claim is
corrected: nothing outranks an as-is pair, but a coarse member whose projected
count equals the physical one ties it, and `maxByOption` would then let
enumeration order decide.
Your suggestion turned up something worth its own paragraph in the
description. Asking "do both children report this pair verbatim" meant
searching each child's partitioning tree, because `joinKeyPositions` could not
answer it. An identity projection keeps every position but still goes through
`toGrouped`, so `createShuffleSpec` reported `Some(0 until n)` even where the
projection and the regrouping rebuilt the very partitioning they started from.
It now reports `None` there, which makes `joinKeyPositions.isEmpty` mean
exactly "this is the child's own layout". The `@param` said "this is set if
joining on a subset of cluster keys is allowed", which was already wrong before
this PR.
That also drops a `GroupPartitionsExec` that projects nothing. The per-child
alignment path matches on `Some(positions)` to decide whether to wrap a child,
so an identity projection used to insert a node whose `grouping` neither
coalesced nor reordered anything. Three existing tests asserted on that node,
and they now assert where the other side's shuffle lands, which is what decides
the plan.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1253,12 +1241,16 @@ case class ShufflePartitionIdPassThrough(
copy(expr = newChildren.head.asInstanceOf[DirectShufflePartitionID])
}
-trait ShuffleSpec {
- /**
- * Returns the number of partitions of this shuffle spec
- */
- def numPartitions: Int
-
+/**
+ * Describes how a child's data is laid out, for the purpose of deciding
whether two children are
+ * co-partitioned and, if not, what to shuffle the other one onto.
+ *
+ * A [[LeafShuffleSpec]] is one concrete layout. A [[ShuffleSpecCollection]]
stands for a choice
+ * between several. A collection can answer [[isCompatibleWith]], which
succeeds when any member
+ * matches. It cannot answer anything that needs one member: which one is
right depends on what the
+ * other side matched, and only the caller comparing the two sides can see
that.
+ */
+sealed trait ShuffleSpec {
Review Comment:
Added to the user-facing section, with your MiMa point: catalyst is in the
`defaultExcludes` section and treated as internals, so CI does not flag it and
no excludes are needed, the same as for the `KeyGrouped*` to `Keyed*` rename.
--
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]
