peter-toth commented on code in PR #58552:
URL: https://github.com/apache/spark/pull/58552#discussion_r3979224148
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala:
##########
@@ -106,9 +106,15 @@ trait ShuffledJoin extends JoinCodegenSupport {
partitionings.map {
case partitioning: Partitioning with Expression
if
PartitioningCollection.keyedMarkerOf(partitioning).contains(true) =>
+ // One cleared layout for the whole input, because a collection's
members must share the
+ // layout by reference. They already share one, so the first
member's answers for all.
+ var cleared: KeyLayout = null
partitioning.transform {
case k: KeyedPartitioning if k.mayContainUnknownPartitionKeys =>
- k.copy(mayContainUnknownPartitionKeys = false)
+ if (cleared == null) {
+ cleared = k.layout.copy(mayContainUnknownPartitionKeys = false)
Review Comment:
Confirmed, and it is a regression this PR introduced. Fixed in
[`72919b8`](https://github.com/apache/spark/commit/72919b80ccaca955b024658cda277d70a1a12648).
The base kept the `partitionKeys` reference through
`k.copy(mayContainUnknownPartitionKeys = false)`, so `intern` returned both
sides untouched. Moving the marker into `KeyLayout` means clearing it has to
build a layout, and a fresh one is `ne` the unmarked side's, so one side
rebuilt on every call. With a nested collection as the mismatching member that
is the whole subtree, and `outputPartitioning` is a `def`.
I took your suggestion. The guard above the block means an unmarked side
always exists, so its layout is looked up once and used as `cleared` when it
equals the copy, with the copy as the fallback. That also let the `var cleared
= null` go, which is your next comment.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -876,19 +919,28 @@ case class KeyedPartitioning(
object KeyedPartitioning {
/**
- * Creates a KeyedPartitioning with isGrouped computed from the partition
keys.
- * Use this when creating a new KeyedPartitioning from scratch (e.g., from a
data source).
+ * Creates a KeyedPartitioning with isGrouped computed from the partition
keys. Use this when
+ * creating a new KeyedPartitioning from scratch (e.g., from a data source).
+ *
+ * `sortKeys` sorts them first, at the types they will be compared at. A
data source reports its
+ * splits in its own order, and a keyed side and a side re-shuffled onto it
have to agree on the
+ * order or `PartitioningCollection.fromPartitionings` refuses them. Sorting
here rather than in
+ * the caller is what keeps the type list and the ordering to one
derivation: both come off the
+ * factory that builds the keys.
*/
def apply(
expressions: Seq[Expression],
- partitionKeys: Seq[InternalRow]): KeyedPartitioning = {
- val dataTypes = expressions.map(_.dataType)
- val comparableKeyWrapperFactory =
-
InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(dataTypes)
- val comparablePartitionKeys =
partitionKeys.map(comparableKeyWrapperFactory)
+ partitionKeys: Seq[InternalRow],
+ sortKeys: Boolean = false): KeyedPartitioning = {
Review Comment:
Taken in
[`72919b8`](https://github.com/apache/spark/commit/72919b80ccaca955b024658cda277d70a1a12648).
`sortKeys` is gone and the scan sorts with
`KeyedPartitioning.groupedKeyRowOrdering(exprs.map(_.dataType))`.
You are right that the "one derivation" property holds either way, and
naming the helper says more than a boolean does: the reader sees that the scan
lays keys out the way `GroupPartitionsExec` and `EnsureRequirements` do. The
`apply` scaladoc now points a caller with unordered keys at that helper instead
of carrying a flag every other caller defaults.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -567,17 +570,24 @@ case class GroupPartitionsExec(
}
/**
- * What a [[GroupPartitionsExec]] computes once and reports from several
members. The last two
- * fields count the alignment's effect on the reads of the child's splits (see
+ * What a [[GroupPartitionsExec]] computes once and reports from several
members: which of the
+ * child's partitions each of its own is built from, and the layout that
describes them. The last
+ * two fields count the alignment's effect on the reads of the child's splits
(see
* `alignToExpectedKeys`), and are 0 outside the alignment path.
*/
private case class PartitionGrouping(
Review Comment:
Taken in
[`72919b8`](https://github.com/apache/spark/commit/72919b80ccaca955b024658cda277d70a1a12648).
`PartitionGrouping` holds a `KeyLayout` in place of `dataTypes`, `isGrouped`
and `isCollapsed`, built where `childKp` is in scope and carrying its marker.
`outputPartitioning` loses the `if (marked) ... else ...` and the comment
that explained the split. It still computes `marked` for the give-up branch,
which is the only thing that reads it now.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ShuffledJoin.scala:
##########
@@ -106,9 +106,15 @@ trait ShuffledJoin extends JoinCodegenSupport {
partitionings.map {
case partitioning: Partitioning with Expression
if
PartitioningCollection.keyedMarkerOf(partitioning).contains(true) =>
+ // One cleared layout for the whole input, because a collection's
members must share the
+ // layout by reference. They already share one, so the first
member's answers for all.
+ var cleared: KeyLayout = null
Review Comment:
Taken in
[`72919b8`](https://github.com/apache/spark/commit/72919b80ccaca955b024658cda277d70a1a12648),
in the shape you wrote, and `representativeOf` widened to `private[sql]` for
it.
The layout is computed once from the representative before the `transform`,
so the guard goes with the `var`. It also gave the fix for your first comment
somewhere to live: the choice between the unmarked side's layout and the fresh
copy is made in that same place.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1229,41 +1275,57 @@ object PartitioningCollection {
representativeOf(p).map(_.mayContainUnknownPartitionKeys)
/**
- * Builds a [[PartitioningCollection]], unifying the `partitionKeys`
reference across all
- * [[KeyedPartitioning]]s (including those in nested collections). Use this
when combining
- * independently-computed partitionings (e.g. join `outputPartitioning`)
where
- * `KeyedPartitioning.partitionKeys` are structurally equal but may not be
reference-equal.
+ * Builds a [[PartitioningCollection]], unifying the [[KeyLayout]] reference
across all
+ * [[KeyedPartitioning]]s, including those in nested collections. Use this
when combining
+ * independently-computed partitionings, such as a join's
`outputPartitioning`, whose layouts
+ * describe the same partitions but are not the same object.
*
* Note: this can't be implemented with `TreeNode.transform`.
*/
def fromPartitionings(partitionings: Seq[Partitioning]):
PartitioningCollection = {
// See the class doc for why the flags are normalized by OR rather than
required to agree. One
- // representative per member is enough, because every collection agrees on
the flags
- // internally by this same construction, and only a member that disagrees
is rebuilt.
+ // representative per member is enough, because every collection agrees
internally by this same
+ // construction, and only a member that disagrees is rebuilt.
val anyCollapsed =
partitionings.exists(representativeOf(_).exists(_.isCollapsed))
val anyUnknownKeys =
partitionings.exists(representativeOf(_).exists(_.mayContainUnknownPartitionKeys))
- var canonicalKeys: Seq[InternalRowComparableWrapper] = null
+ var canonicalLayout: KeyLayout = null
// A partitioning with no `KeyedPartitioning` in it has nothing to
normalize, and one that
- // already agrees on the keys and both flags is returned as it is. That is
what keeps
- // repeated `outputPartitioning` computations over deeply nested
collections (e.g. chains of
- // same-key joins) O(1) per level.
+ // already holds the canonical layout is returned as it is. That is what
keeps repeated
+ // `outputPartitioning` computations over deeply nested collections (e.g.
chains of same-key
+ // joins) O(1) per level.
def intern(p: Partitioning): Partitioning = representativeOf(p) match {
case None => p
case Some(representative) =>
- if (canonicalKeys == null) canonicalKeys = representative.partitionKeys
- if ((representative.partitionKeys eq canonicalKeys) &&
- representative.isCollapsed == anyCollapsed &&
- representative.mayContainUnknownPartitionKeys == anyUnknownKeys) {
+ if (canonicalLayout == null) {
+ val layout = representative.layout
+ canonicalLayout =
+ if (layout.isCollapsed == anyCollapsed &&
+ layout.mayContainUnknownPartitionKeys == anyUnknownKeys) {
+ layout
+ } else {
+ layout.copy(
+ isCollapsed = anyCollapsed, mayContainUnknownPartitionKeys =
anyUnknownKeys)
+ }
+ }
+ if (representative.layout eq canonicalLayout) {
p
} else {
- require(representative.partitionKeys == canonicalKeys,
+ require(representative.partitionKeys ==
canonicalLayout.partitionKeys,
"All KeyedPartitionings in a PartitioningCollection must have
equal partitionKeys")
+ // Whether the keys are unique is a property of the keys, so two
layouts over equal keys
+ // that disagree on it cannot both be right.
+ require(representative.isGrouped == canonicalLayout.isGrouped,
+ "All KeyedPartitionings in a PartitioningCollection must agree on
isGrouped")
+ // Interning replaces a member's layout whole, so a member that
describes another key
+ // space would be silently retyped. Two empty key lists compare
equal whatever they
+ // describe, which is the case the clause above them cannot see.
+ require(representative.keyDataTypes == canonicalLayout.dataTypes,
Review Comment:
Taken in
[`72919b8`](https://github.com/apache/spark/commit/72919b80ccaca955b024658cda277d70a1a12648).
`KeyLayout.describesSameKeys` holds the predicate, and both `intern` and
`KeyedShuffleSpec.isCompatibleWith` call it.
The rationale moved with it: why the types are asked as well as the rows,
and the residual you and ulysses-you both landed on, that two empty sides
sharing a type pair anyway and why that is inert. `isGrouped` stays a separate
`require`, and the scaladoc says why rather than leaving the next reader to
guess.
The `keyDataTypes` versus `dataTypes` mixture is gone from the requires too.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1941,16 +2018,18 @@ case class KeyedShuffleSpec(
te.copy(children = te.children.map(_ => clustering(positionSet.head)))
case (_, positionSet) => clustering(positionSet.head)
}
- // The shuffled side is laid out on this side's partition keys, so it
inherits the flag. That
- // is conservative rather than strictly true, and it can only ever add a
shuffle: a later
- // grouping of the shared key set carries the collapsed side's risk.
+ // The shuffled side is laid out on this side's partitions, so it shares
their layout, with one
+ // change. The child re-shuffled onto it may hold keys outside the
declared set, so every
+ // partitioning produced here carries the marker (see [[KeyLayout]]'s
`@param`). Only the
+ // shuffle loop reaches this call, and it carries no same-domain subset
proof, so marking is
+ // sound; it is conservative where the child's keys are in fact a known
subset (identity key
+ // [1] inside declared [1, 2]), a precision this path does not attempt.
//
- // The child re-shuffled onto this layout may hold keys outside the
declared set, so every
- // partitioning produced here carries the marker (see the `@param`). Only
the shuffle loop
- // reaches this call, and it carries no same-domain subset proof, so
marking is sound; it is
- // conservative where the child's keys are in fact a known subset
(identity key [1] inside
- // declared [1, 2]), a precision this path does not attempt.
- partitioning.copy(expressions = newExpressions,
mayContainUnknownPartitionKeys = true)
+ // There is nothing to decide about `isCollapsed`: a later grouping of the
shared key set
+ // carries the same risk whichever side reports it.
+ partitioning
+ .copy(expressions = newExpressions)
+ .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
Review Comment:
Taken in
[`72919b8`](https://github.com/apache/spark/commit/72919b80ccaca955b024658cda277d70a1a12648).
One `copy` with both arguments, no intermediate node.
--
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]
