dongjoon-hyun commented on code in PR #58351:
URL: https://github.com/apache/spark/pull/58351#discussion_r3881703132
##########
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:
`projectedDistinctKeyCount` is taken from `keyToPartitionIndices.size`
*before* `alignToExpectedKeys` prunes to the join-agreed key set, so a collapse
confined entirely to keys that the intersection then filters out still marks
the output collapsed, even though every surviving partition maps 1:1 to a
source partition.
Reachable on the reducer path with `allowCompatibleTransforms` +
`partitionFilter` on and the subset opt-in off: ids `[0, 4, 5]` reduced onto
buckets `[0, 0, 1]` with bucket 0 pruned by the other side gives `2 < 3`,
flagging an output whose only surviving partition was never merged. Since the
flag is sticky, a later union + GROUP BY then hits the `isCollapsed &&
!isGrouped` gate and pays a shuffle the pre-PR code planned shuffle-free.
The new "filtering partition keys out is not a key collapse" test covers
pruning *without* a collapse, but not this variant where the pruned keys are
themselves the collapsed ones. Perf-only, never wrong results, but it
contradicts the flag's documented meaning ("one partition here can stand for
several of the original ones").
##########
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:
`outputPartitioning` is a `def`, so this block re-runs on every call:
`flatten` allocates a Seq just to `collectFirst` one KP, and
`k.distinctKeyCount` is forced on a fresh child KP instance each time
(`DataSourceV2ScanExecBase.outputPartitioning` rebuilds its `KeyedPartitioning`
per call), so the lazy val never amortizes -- an O(#splits) distinct pass per
consultation during EnsureRequirements/validation/AQE.
Computing the flag once inside the memoized `groupedPartitionsTuple` --
which already `collectFirst`s the same child KP and holds
`keyToPartitionIndices.size` -- is semantically identical, and would also
remove the unreachable `.getOrElse(false)` (`groupedPartitions`, forced just
above via `partitionKeys`, throws when no KP exists) and the second, divergent
first-KP lookup idiom in this file (line 148 uses plain `collectFirst`).
##########
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:
Since `isCollapsed` is a constructor param, it participates in case-class
equality and survives canonicalization. Two otherwise-identical exchanges --
one templated from a collapsed partitioning, one from an equivalent
non-collapsed one (two independent joins over a shared subplan; the collection
OR-normalization doesn't reach across them) -- no longer compare equal under
`sameResult`, so `ReuseExchangeAndSubquery` shuffles the shared subplan twice
where it previously reused the exchange.
Narrow configuration, and arguably defensible since the metadata genuinely
differs (reusing a false-flag exchange where a true one was planned would
re-launder the protection under AQE re-planning) -- but it deserves a conscious
decision.
##########
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:
The inherited flag is also stamped on join types that expose only the
shuffled side (`LeftOuter`/`LeftSemi`/`LeftAnti`/`LeftExistence` in
`ShuffledJoin.outputPartitioning`); for semi/anti the output carries zero
collapsed-side rows, yet inherits the flag forever. The created KP is grouped
(`canCreatePartitioning` requires `isGrouped`), so the flag only bites after a
union reintroduces duplicate keys -- exactly the case the class doc blesses as
groupable without opt-in -- and the flag's documented contract is false for
these partitionings.
The spec can't know the consuming join type, so if this is worth tightening,
the place would be `ShuffledJoin.outputPartitioning` (clear the flag for
`LeftExistence`). Fine to keep as the documented safe-direction conservatism,
but worth noting the cost and the labeling mismatch.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala:
##########
@@ -137,12 +137,24 @@ trait PartitioningPreservingUnaryExecNode extends
UnaryExecNode
if (projectablePositions.length == numPositions) keySource.partitionKeys
else keySource.projectKeys(projectablePositions)._2
- val isGrouped = sharedKeys.distinct.size == sharedKeys.size
- // A KP is narrowed if this node drops positions, or if the input KPs were
already narrowed
- // (i.e. came from a finer-grained partitioning). The flag must be sticky:
a subsequent
- // PartitioningPreservingUnaryExecNode that passes all positions through
would otherwise
- // recompute isNarrowed=false, silently dropping the protection.
- val isNarrowed = projectablePositions.length < numPositions ||
keySource.isNarrowed
+ val distinctSharedKeys = sharedKeys.distinct
+ val isGrouped = distinctSharedKeys.size == sharedKeys.size
+ // This projection collapses keys when it maps keys that were distinct in
the input
+ // onto the same projected key -- dropping positions is not enough on its
own, since the
+ // projected keys can stay just as distinct as the originals. The flag is
sticky: a subsequent
+ // PartitioningPreservingUnaryExecNode that passes all positions through
must not recompute it
+ // as false and drop the protection, and no projection can make a
partitioning finer again.
+ //
+ // Both cheap terms come first: an inherited flag or a projection that
drops no position
+ // settles the question without counting distinct keys. A pass-through
projection cannot
+ // collapse anything, since it keeps the input's keys as they are.
+ //
+ // The inherited flag is read from all inputs rather than from the key
source alone. A
+ // `PartitioningCollection` normalizes it across its members, so the two
agree today; reading
+ // all of them keeps this producer correct without depending on that.
+ val isCollapsed = kps.exists(_.isCollapsed) ||
Review Comment:
This `inherited || projected distinct < source distinct` predicate is now
hand-rolled at three producer sites: here,
`KeyedPartitioning.createShuffleSpec`, and
`GroupPartitionsExec.outputPartitioning`, each with site-specific inputs and
caveats. Given the PR's own observation that one producer laundering the flag
is how the protection went missing, a shared helper on `KeyedPartitioning`, e.g.
```scala
def collapsedAfterProjection(projectedDistinctKeyCount: Int): Boolean =
isCollapsed || projectedDistinctKeyCount < distinctKeyCount
```
would keep the correctness rule in one place instead of three files that
must stay in sync.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -643,12 +676,17 @@ case class KeyedPartitioning(
val joinKeyPositions =
result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2)
val projectedExpressions = joinKeyPositions.map(expressions)
val projectedKeys = projectKeys(joinKeyPositions)._2
+ // Projecting onto the operation keys can collapse keys in its own
right, which is
+ // what the key count comparison catches. The gate in `groupedSatisfies`
is bypassed while
+ // this config is on, so the flag decides nothing here today, but it
travels with the
+ // partitioning and leaving a producer to launder it is how the
protection went missing.
+ val projectedCollapsed = isCollapsed || projectedKeys.distinct.length <
distinctKeyCount
Review Comment:
Two O(n) distinct passes over the same list: `projectedKeys.distinct.length`
here, then `.toGrouped` on the next statement re-runs
`partitionKeys.distinct.sorted` over the same `projectedKeys`; for an
ungrouped, uncollapsed source, forcing `distinctKeyCount` adds a third pass
over the source keys. The block also runs when `joinKeyPositions` selects every
position, where the comparison is tautologically false.
Materializing `val d = projectedKeys.distinct` once and building the grouped
KP directly (same ordering source as `toGrouped`), and/or guarding on positions
actually dropped, is semantically identical and cheaper. Planning-time only, so
minor -- but it runs per `createShuffleSpec` call from
EnsureRequirements/ValidateRequirements.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -528,22 +554,25 @@ case class CoalescedNullAwareHashPartitioning(
* guaranteed after projection. May contain duplicates
when ungrouped.
* @param isGrouped Whether partition keys are unique (no duplicates).
Computed on first
* creation, then preserved through copy operations to avoid
recomputation.
- * @param isNarrowed Whether this partitioning was derived from a
finer-grained one by dropping key
- * positions (e.g. via
`PartitioningPreservingUnaryExecNode`). When true and the
- * keys are no longer unique, `GroupPartitionsExec` may
merge partitions that held
- * distinct keys in the original partitioning, carrying the
same skew risk as
- * `allowKeysSubsetOfPartitionKeys`. "May", because the
condition is a proxy: the
- * duplicate keys can also come from a source that reports
several splits per
- * partition key, in which case grouping merges only
same-key partitions. Such a
- * partitioning can only satisfy `ClusteredDistribution` by
being grouped, and
- * `groupedSatisfies` refuses that unless the config is
enabled, regardless of
- * `requireAllClusterKeysForDistribution`.
+ * @param isCollapsed Whether a projection or a reduction mapped keys that
were distinct in the
+ * partitioning this one was derived from onto the same
key, so one partition
+ * here can stand for several of the original ones -- see
"Key Collapse" above.
+ * Dropping key positions does not set it on its own; the
projected keys have to
+ * actually lose distinctness. Sticky, because neither
grouping nor a further
+ * projection can make a partitioning finer again. One case
sets it without a
+ * collapse of its own: the side shuffled onto a collapsed
partitioning's keys
+ * inherits it, because the two are then co-located on that
key set -- see
+ * `KeyedShuffleSpec.createPartitioning`.
+ * Together with `!isGrouped` it decides whether
`groupedSatisfies` may coalesce
+ * the duplicate keys without
`allowKeysSubsetOfPartitionKeys`: `isCollapsed`
+ * says the collapse happened, `!isGrouped` says there is
still something left
+ * to merge, and only both together mean there is an
outstanding risk to gate.
*/
case class KeyedPartitioning(
expressions: Seq[Expression],
@transient partitionKeys: Seq[InternalRowComparableWrapper],
isGrouped: Boolean,
- isNarrowed: Boolean = false) extends Expression with Partitioning with
Unevaluable {
+ isCollapsed: Boolean = false) extends Expression with Partitioning with
Unevaluable {
Review Comment:
Keeping `isCollapsed: Boolean = false` as a defaulted parameter leaves the
laundering vector this PR closes open for the next producer: a future operator
that derives a KeyedPartitioning and calls the 3-arg constructor silently gets
`false` and bypasses the gate with no compiler or test signal -- the same
failure mode as SPARK-59026, and the wrong code looks identical to correct
pre-existing code.
Only the companion `apply` (a provably-fresh source partitioning) actually
relies on the default; dropping the default and passing `isCollapsed = false`
explicitly there would make the compiler force every future producer to decide
the flag.
--
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]
