peter-toth commented on code in PR #58351:
URL: https://github.com/apache/spark/pull/58351#discussion_r3882330517
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -68,13 +68,27 @@ case class GroupPartitionsExec(
child.outputPartitioning match {
case p: Partitioning with Expression =>
// There can be multiple `KeyedPartitioning`s in an output
partitioning of a join, but they
- // can only differ in `expressions`; their `partitionKeys` reference
is shared (enforced by
- // `PartitioningCollection`), so `groupedPartitions` is computed only
once.
+ // can only differ in `expressions`; their `partitionKeys` reference
and `isCollapsed` flag
+ // are shared (enforced by `PartitioningCollection`), so both
`groupedPartitions` and the
+ // new flag are computed once, outside the transform.
val partitionKeys = groupedPartitions.map(_._1)
+ // `isCollapsed` is sticky: grouping removes the duplicate keys, but
it does not make this
+ // partitioning any finer than the layout it came from. Projecting
onto the join key
+ // positions, or reducing the keys onto a coarser transform, can
collapse keys in its own
+ // right, which is what the key count comparison catches. Compare
against this side's
+ // own key count, not `partitionKeys`: that list is the one both join
sides agreed on, so it
+ // may be missing keys this side had (partition filtering) or repeat
them (padding).
+ // One member is enough: they share the `partitionKeys` reference and
the flag, so they
+ // also share `distinctKeyCount`. Reading them all would force that
count -- a pass over
+ // the keys -- once per member for the same answer.
+ val isCollapsed = PartitioningCollection.flatten(p).collectFirst {
+ case k: KeyedPartitioning =>
+ k.isCollapsed || projectedDistinctKeyCount < k.distinctKeyCount
Review Comment:
Good catch, and it took two goes to get right.
The flag no longer compares key counts at all. `groupedPartitionsTuple` now
asks the question directly, of the key groups it keeps: does any of them cover
more than one of the child's own partition keys? Counting the child's *keys*
rather than its partitions is what tells a collapse from a source reporting
several splits per key, and asking it of the kept groups is what tells it from
filtering. Your shape is a test now: `identity(id)` with ids 0, 4, 5 reduced
onto buckets `[0, 0, 1]`, bucket 0 dropped because the other side has no rows
for it, so the only surviving key covers id 5 alone.
The second go: my first version read the answer off the partitions
`alignToExpectedKeys` emits, which is blind under `distributePartitions` --
that branch spreads a group's splits over one partition each, so no partition
ever holds two, and the whole mode reported nothing collapsed. That was a
regression against the state you reviewed. Asking the key groups is right for
both modes, and there is a second test for that direction, so the two tests are
each other's control.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -68,13 +68,27 @@ case class GroupPartitionsExec(
child.outputPartitioning match {
case p: Partitioning with Expression =>
// There can be multiple `KeyedPartitioning`s in an output
partitioning of a join, but they
- // can only differ in `expressions`; their `partitionKeys` reference
is shared (enforced by
- // `PartitioningCollection`), so `groupedPartitions` is computed only
once.
+ // can only differ in `expressions`; their `partitionKeys` reference
and `isCollapsed` flag
+ // are shared (enforced by `PartitioningCollection`), so both
`groupedPartitions` and the
+ // new flag are computed once, outside the transform.
val partitionKeys = groupedPartitions.map(_._1)
+ // `isCollapsed` is sticky: grouping removes the duplicate keys, but
it does not make this
+ // partitioning any finer than the layout it came from. Projecting
onto the join key
+ // positions, or reducing the keys onto a coarser transform, can
collapse keys in its own
+ // right, which is what the key count comparison catches. Compare
against this side's
+ // own key count, not `partitionKeys`: that list is the one both join
sides agreed on, so it
+ // may be missing keys this side had (partition filtering) or repeat
them (padding).
+ // One member is enough: they share the `partitionKeys` reference and
the flag, so they
+ // also share `distinctKeyCount`. Reading them all would force that
count -- a pass over
+ // the keys -- once per member for the same answer.
+ val isCollapsed = PartitioningCollection.flatten(p).collectFirst {
Review Comment:
Done. The flag is computed inside `groupedPartitionsTuple`, so
`outputPartitioning` no longer forces `distinctKeyCount` on a child
partitioning that `DataSourceV2ScanExecBase` rebuilt on every call. The child
lookup is now a shared `childKeyedPartitioning` used by both the grouping and
the flag, which removes the second idiom and the unreachable `getOrElse(false)`
as well.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1445,7 +1513,13 @@ case class KeyedShuffleSpec(
te.copy(children = te.children.map(_ => clustering(positionSet.head)))
case (_, positionSet) => clustering(positionSet.head)
}
- KeyedPartitioning(newExpressions, partitioning.partitionKeys,
partitioning.isGrouped)
+ // The shuffled side is laid out on this side's partition keys, so it
inherits the flag.
+ // Strictly nothing collapsed on this side -- its partitions are what a
hash partitioning would
+ // give -- so this is deliberate conservatism: the two sides are
co-located on one key set, and
+ // a later grouping of that key set carries the collapsed side's risk. It
can only add shuffles,
+ // never remove one.
+ KeyedPartitioning(newExpressions, partitioning.partitionKeys,
partitioning.isGrouped,
Review Comment:
Conscious decision: keep it. The metadata genuinely differs, and reusing an
exchange planned from a non-collapsed partitioning where a collapsed one was
expected would re-launder the protection under AQE re-planning, which is the
failure this change exists to close. The cost needs the opt-in plus two
independent joins over a shared subplan, and it is a lost reuse rather than a
wrong answer.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1445,7 +1513,13 @@ case class KeyedShuffleSpec(
te.copy(children = te.children.map(_ => clustering(positionSet.head)))
case (_, positionSet) => clustering(positionSet.head)
}
- KeyedPartitioning(newExpressions, partitioning.partitionKeys,
partitioning.isGrouped)
+ // The shuffled side is laid out on this side's partition keys, so it
inherits the flag.
+ // Strictly nothing collapsed on this side -- its partitions are what a
hash partitioning would
+ // give -- so this is deliberate conservatism: the two sides are
co-located on one key set, and
+ // a later grouping of that key set carries the collapsed side's risk. It
can only add shuffles,
+ // never remove one.
+ KeyedPartitioning(newExpressions, partitioning.partitionKeys,
partitioning.isGrouped,
+ partitioning.isCollapsed)
Review Comment:
I would keep it, and I do not think the labelling is wrong. The flag
describes key granularity, not which rows survive: a semi-join output laid out
on the template's collapsed key space really does have partitions covering
several of the source's keys, whichever side's rows flow through. Clearing it
for `LeftExistence` would claim the layout is finer than it is, and the first
union that reintroduces duplicate keys would then group genuinely coarse
partitions without the opt-in.
--
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]
