gengliangwang commented on code in PR #58245:
URL: https://github.com/apache/spark/pull/58245#discussion_r3847107324
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -61,6 +61,10 @@ case class EnsureRequirements(
shuffleOrigin: ShuffleOrigin): Seq[SparkPlan] = {
assert(requiredChildDistributions.length == originalChildren.length)
assert(requiredChildOrderings.length == originalChildren.length)
+ // A storage-partitioned join handles its co-partitioning (and the
projected join-key
+ // GroupPartitionsExec) separately in checkKeyGroupCompatible, so the
projected-key grouping
+ // below must not run for it. For a non-join operator, do it inline here.
+ val isJoin = parent.exists(_.isInstanceOf[ShuffledJoin])
Review Comment:
This guard states a different condition from the one the comment appeals to,
and the two sets come apart in both directions.
The comment says joins are excluded because `checkKeyGroupCompatible`
handles their projection. But that helper matches only `SortMergeJoinExec` and
`ShuffledHashJoinExec`, while the block that calls it is entered for *any*
operator with two clustered children (`parent.isDefined && children.length == 2
&& childrenIndexes.length == 2`).
One direction is currently harmless: `SortMergeAsOfJoinExec` is also a
`ShuffledJoin`, so it is excluded here yet unhandled there. It still gets
correct results, but by falling through to a plain shuffle rather than by
anything this comment describes.
The other direction concerns me, and I could not finish confirming it -
hence a question. `CoGroupExec` requires `ClusteredDistribution` on both
children (objects.scala:638) and is not a `ShuffledJoin`, so this branch does
run for it. `children` is then reassigned with the wrapped child, and `specs` a
few lines down are computed from `children(i).outputPartitioning` - the
*projected* partitioning. `checkKeyGroupCompatible` returns `None`, so
`areChildrenCompatible` is false and each clustered child reaches
`withJoinKeyPositions(child, joinKeyPositions)`, which rewrites the node in
place via `g.copy(joinKeyPositions = Some(positions))`. Those positions index
the projected expression list, but `GroupPartitionsExec` applies them to its
child's unprojected partitioning. For tables partitioned by `(name, id)`
cogrouped on `id`: this branch computes `[1]` and inserts a node projecting to
`id`; the block recomputes `[0]` against the now one-element list and
overwrites, so the node projec
ts position 0 of `(name, id)` and coalesces by `name`.
I did not verify the observable result. Reaching the overwrite also needs
`v2BucketingShuffleEnabled` on, since it gates
`KeyedShuffleSpec.canCreatePartitioning` - with the default false,
`bestSpecOpt` is empty and the fallback shuffles and unwraps the node instead.
Could you confirm whether a cogroup over two storage-partitioned tables with
both configs on and a non-leading grouping key produces wrong groups?
Either way I'd gate on the structural fact rather than the trait: run this
branch only when the operator will not enter that block, i.e.
`childrenIndexes.length <= 1`. That is what the comment is really appealing to,
it covers cogroups and as-of joins without enumerating join classes, and it
keeps one owner of the projection per operator so a future `ShuffledJoin` or a
new parent case in `checkKeyGroupCompatible` cannot move the boundary silently.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4207,6 +4207,111 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
+ test("window top-k over PARTITION BY subset of partition keys coalesces
partitions") {
Review Comment:
Consider adding an aggregate variant next to these. `HashAggregateExec`
requires `ClusteredDistribution(groupingExpressions)`, so `SELECT id,
sum(price) FROM items GROUP BY id` on an `(id, name)`-partitioned table reaches
this exact branch, and without the fix it emits one row per `(id, name)`
partition instead of per `id` - a duplicated group, which is arguably a more
familiar symptom than a ranking artifact. All three new tests go through
`ROW_NUMBER() OVER (PARTITION BY ...)`, so a reader could reasonably think the
fix is window-specific.
A cogroup test would be worth more still, because that case is one the
change newly *affects* rather than fixes:
`df1.groupByKey(...).cogroup(df2.groupByKey(...))` over two storage-partitioned
tables is a two-clustered-child operator that is not a `ShuffledJoin`, so it
takes the new branch and then also goes through the co-partitioning block. That
is the interaction I asked about on the `isJoin` line.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -107,8 +111,24 @@ case class EnsureRequirements(
}
case _ if groupedSatisfies.isDefined =>
- // Grouped KeyedPartitioning already satisfies
- child
+ // A grouped KeyedPartitioning already satisfies. However,
when the operation keys
+ // are a strict subset of the partition keys (enabled via
+ // v2BucketingAllowKeysSubsetOfPartitionKeys), the partitions
are still grouped by
+ // the full partition keys rather than by the operation keys,
so a
+ // GroupPartitionsExec that projects to the operation keys
must be inserted to
+ // coalesce partitions sharing the same operation key.
`createShuffleSpec` computes
+ // exactly those projected positions when the config is
enabled.
+ val kp = groupedSatisfies.get
+ distribution match {
+ case c: ClusteredDistribution if !isJoin =>
+ val spec =
kp.createShuffleSpec(c).asInstanceOf[KeyedShuffleSpec]
Review Comment:
Only `spec.joinKeyPositions` is used, but with the config enabled
`createShuffleSpec` also runs `projectKeys(joinKeyPositions)._2` across every
partition key and then `.distinct` on the result, to build a
`projectedPartitioning` that this call site drops - and `GroupPartitionsExec`
recomputes the same projection later from `child.outputPartitioning`. That is
two O(number of partitions) passes per qualifying operator at planning time, on
a path that exists precisely for tables partitioned finely enough to need
coalescing.
`KeyedShuffleSpec(kp, c).keyPositions` gives the same information without
the discarded projection, and reads more directly as "which partition
expressions does this operation actually key on" than routing through a
shuffle-spec factory. The cast also goes away with it.
--
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]
