LuciferYang commented on code in PR #58531:
URL: https://github.com/apache/spark/pull/58531#discussion_r3945718698
##########
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:
The PR description's argument that the as-is pair "can never be lost" only
holds under strict inequality. Projected counts never exceed physical ones, so
no pair strictly outranks the as-is pair, but `maxByOption` breaks ties by
enumeration order, and a coarse member whose projected distinct count happens
to equal the physical count, listed before the identity member, ties it and
wins. Each side then gets an extra `GroupPartitionsExec`, and the grouping
drops the other key the identity member still carried. This is not a regression
against the old code, which picked the same head members, so the minimal fix is
correcting the claim in the description. If you want the cheap improvement,
break ties in favor of the pair whose specs' partitionings appear verbatim in
both children.
##########
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:
The `ShuffleSpec` family is public in 4.0.0, and this commit removes
members, adds an abstract `flatten`, and seals the trait. Downstream code that
still extends catalyst internals sees a real source and binary break: `extends
ShuffleSpec` no longer compiles, and old bytecode calling `numPartitions` hits
NoSuchMethodError. The whole catalyst package is in the defaultExcludes section
of `MimaExcludes` and treated as internals, though, so CI will not flag it and
no excludes are needed; the KeyGrouped* to Keyed* rename set that precedent.
I'd still add a line to the user-facing-change section of the PR
description: `sealed` is a source break for external `extends ShuffleSpec`
code, which needs to move to `LeafShuffleSpec`.
##########
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:
The pair ranking `l.numPartitions.max(r.numPartitions)` measures one side's
projected count, but what the push branch delivers is the merged key set. Under
partition filter + Inner that is the intersection, and its size is not monotone
in member granularity: a strictly finer pair can win the ranking while its
intersection with the other side comes out far smaller than the coarse pair's,
so the join ends up with fewer tasks than the old head-of-list pick produced.
Ranking by the key count the pair actually merges would match the delivered
parallelism, but it means computing reduced keys for every candidate pair,
which is expensive and can surface reduced-types mismatches at ranking time for
pairs that would never be chosen. At minimum, the comment should say what `max`
measures and where it stops matching the delivered parallelism.
--
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]
