dongjoon-hyun commented on code in PR #58351:
URL: https://github.com/apache/spark/pull/58351#discussion_r3882512619
##########
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:
Nit, non-blocking: this copies even when `projectedCollapsed ==
grouped.isCollapsed` (both false in the common no-collapse case); `if
(projectedCollapsed) grouped.copy(isCollapsed = true) else grouped` avoids the
redundant allocation per spec creation.
##########
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:
Nit, non-blocking: this re-derives the kept groups with its own
`keyToPartitionIndices.getOrElse` lookups, duplicating the key matching
`alignToExpectedKeys` just performed -- if that method's matching ever changes
(normalization, emitting instead of dropping unexpected keys), the emitted
partitions and the flag would silently be computed from different group
selections. Computing the collapse bit inside `alignToExpectedKeys` in the same
pass would keep one source of truth; short of that, fusing the map into
`expectedPartitionKeys.get.exists { case (key, _) =>
coversSeveralChildKeys(keyToPartitionIndices.getOrElse(key, Seq.empty)) }` at
least drops the intermediate Seq and short-circuits.
--
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]
