peter-toth commented on code in PR #58351:
URL: https://github.com/apache/spark/pull/58351#discussion_r3887557545
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -145,10 +158,26 @@ case class GroupPartitionsExec(
val keyToPartitionIndices = reducedKeys.zipWithIndex.groupMap(_._1)(_._2)
+ // Whether this node collapses keys: does any key it keeps stand for 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 keys
+ // this node keeps is what tells it from `alignToExpectedKeys` dropping
keys, which merges
+ // nothing. Ask the key groups, not the partitions finally emitted:
`distributePartitions`
+ // spreads a group's splits over one partition each, which would hide the
merge, and
+ // replication would ask about the same group repeatedly.
+ val childKeys = keyedPartitioning.partitionKeys.toIndexedSeq
+ def coversSeveralChildKeys(indices: Seq[Int]): Boolean =
+ indices.map(childKeys).distinct.size > 1
+
if (expectedPartitionKeys.isDefined) {
- alignToExpectedKeys(keyToPartitionIndices)
+ val (alignedPartitions, grouped) =
alignToExpectedKeys(keyToPartitionIndices)
+ val keptGroups = expectedPartitionKeys.get.map { case (key, _) =>
Review Comment:
This is your second form now, the `exists` over `expectedPartitionKeys`, so
there is no intermediate Seq and it short-circuits.
I went your first way to begin with and computed the collapse bit inside
`alignToExpectedKeys`, then moved it back out. Inside, the aligner had to grow
a mutable buffer and a third tuple element that it never used itself, and its
doc had to explain a return value that belonged to the caller.
What the current form gives up is the single lookup site you were after.
`keyToPartitionIndices.get(key)` now appears twice in the file, and if the
aligner's key matching changed, the two could drift. One thing did get better
with the move: both branches of the scan take their groups from
`keyToPartitionIndices`, so a padded expected key cannot produce an empty group
at all, and the `getOrElse(key, Seq.empty)` plus its guard are gone.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -643,12 +686,23 @@ case class KeyedPartitioning(
val joinKeyPositions =
result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2)
val projectedExpressions = joinKeyPositions.map(expressions)
val projectedKeys = projectKeys(joinKeyPositions)._2
- // Sort the distinct projected keys the same way `GroupPartitionsExec`
does (both sort with
- // `KeyedPartitioning.groupedKeyRowOrdering`). Otherwise, when only the
keyed side is grouped
- // and the other side is re-shuffled using this spec, the two
`KeyedPartitioning`s carry the
- // same keys in a different order and
`PartitioningCollection.fromPartitionings` rejects them.
- val projectedPartitioning =
- new KeyedPartitioning(projectedExpressions, projectedKeys, isGrouped =
false).toGrouped
+ // `toGrouped` dedups and sorts the keys the same way
`GroupPartitionsExec` does (both sort
+ // with `KeyedPartitioning.groupedKeyRowOrdering`). Otherwise, when only
the keyed side is
+ // grouped and the other side is re-shuffled using this spec, the two
`KeyedPartitioning`s
+ // carry the same keys in a different order and
`PartitioningCollection.fromPartitionings`
+ // rejects them. Its `distinct` is also the only one needed here: the
partition count it
+ // leaves is the projected distinct key count the collapse test asks for.
+ val grouped = new KeyedPartitioning(
+ projectedExpressions, projectedKeys, isGrouped = false, isCollapsed =
false).toGrouped
+ // Projecting onto the operation keys can collapse keys in its own
right. Dropping no position
+ // cannot, so the counts are only compared when one was dropped. 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 ||
+ (joinKeyPositions.length < expressions.length &&
+ collapsesOnProjection(grouped.numPartitions))
+ val projectedPartitioning = grouped.copy(isCollapsed =
projectedCollapsed)
Review Comment:
Moot now. `createShuffleSpec` does not copy at all any more, it is
`project(joinKeyPositions).toGrouped`, and `project` decides the flag in the
constructor call.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -491,14 +491,40 @@ case class CoalescedNullAwareHashPartitioning(
* `ClusteredDistribution`, this is the only route by which a grouped KP
satisfies one, and no
* grouping is involved: the keys are already unique.
*
- * That second caller is why the narrowing guard in `groupedSatisfies()` is a
conjunction with
- * `!isGrouped`. A narrowed KP whose projected keys stayed distinct is
grouped, and dropping the
- * `!isGrouped` term would stop it from satisfying a `ClusteredDistribution`
and cost it a shuffle,
- * even though grouping it would merge nothing.
+ * That second caller is why the collapse guard in `groupedSatisfies()` is a
conjunction with
+ * `!isGrouped`. A collapsed KP can end up grouped -- by
`GroupPartitionsExec`, or by reducing its
+ * keys onto a coarser transform -- and dropping the `!isGrouped` term would
stop such a KP from
+ * satisfying a `ClusteredDistribution` and cost it a shuffle, even though
grouping it would merge
+ * nothing.
*
* For `OrderedDistribution`, `GroupPartitionsExec` must also sort the
partition keys to meet the
* ordering requirement.
*
+ * == Key Collapse ==
+ * Two things happen to partition keys, and only the first is a loss of
granularity:
+ *
+ * - '''Key collapse''': a projection or a reduction maps two keys that were
distinct onto the same
+ * new key. `[(1, 'a'), (1, 'b'), (2, 'c')]` projected onto the first
position gives `[1, 1, 2]`:
+ * three distinct keys became two. `isCollapsed` records this.
+ * - '''Grouping''': `GroupPartitionsExec` physically combines the partitions
that share a key.
+ * `[1, 1, 2]` becomes `[1, 2]`. `isGrouped` says the keys are unique,
however they got that way:
+ * a source with natively unique keys reports it too.
+ *
+ * Grouping after a collapse is what produces a partition holding more data
than any the source
+ * declared -- the two `1` partitions above came from different `(1, 'a')` and
`(1, 'b')` keys -- so
+ * it needs `allowKeysSubsetOfPartitionKeys`. Grouping without a collapse only
merges partitions
+ * that already shared a key (a source reporting several splits per key, or a
union of children that
+ * overlap), and needs no opt-in. `OrderedDistribution` is not gated at all:
`GroupPartitionsExec`
+ * pads that path out to the expected split counts rather than coalescing, so
nothing is merged.
+ *
+ * A collapsed partitioning is still kept rather than dropped to
`UnknownPartitioning`, because
Review Comment:
Done, in your words:
- **Collapsed and ungrouped**: duplicate keys remain, so grouping them would
merge partitions the source held apart. `mayGroupToSatisfy()` refuses a
`ClusteredDistribution` unless the config is on.
- **Collapsed and grouped**: the keys are already unique, so grouping merges
nothing and `satisfies()` accepts a `ClusteredDistribution` whatever the config
says. A partitioning reaches this state by being grouped with the config on, or
by having its keys reduced onto a coarser transform.
The paragraph that mixed the two states is gone. What follows the bullets is
only what is common to both, which is why a collapsed partitioning is still
reported rather than dropped to `UnknownPartitioning`.
##########
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:
Update after a rewrite, because the helper I named here no longer exists.
`collapsesOnProjection` is gone, and the rule it held moved inside
`KeyedPartitioning.project`, which both producers now call for the whole
projection. So the rule is in one place rather than the three you started from,
and there is no helper left that a future producer could forget to call.
`GroupPartitionsExec` still answers the question its own way, from the key
groups it keeps, for the reason in my earlier reply.
##########
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:
Update, and the intermediate state was worse than what I told you here.
Extracting the projection into `KeyedPartitioning.project` brought the second
`distinct` back for a while, because `project` counted the projected keys and
`toGrouped` then deduped the same list again.
It is one pass now, and not a `distinct` at all. `project` walks the
projected keys alongside the keys they came from and fills a `projectedKey ->
sourceKey` map. A second, different source key on an existing projected key is
the collapse, and it also means the projected keys are not unique, so the walk
stops there. The source's own distinct count is not computed any more, and the
source keys are never hashed, only compared where a projected key repeats.
Measured on the worst case for that test, a 50k-split partitioning with 25k
distinct keys and 12-position keys, projected ten times over, 20 evaluations:
345 ms for the old provenance formula, 1291 ms for this one, 2137 ms for the
two-`distinct` form. `toGrouped` after it also skips its own dedup when the
keys are already unique.
--
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]
