dongjoon-hyun commented on code in PR #58552:
URL: https://github.com/apache/spark/pull/58552#discussion_r3976989647
##########
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:
Small planning-time regression here. `cleared` is a fresh `KeyLayout`, so it
is never `eq` the keyed side's layout, and `fromPartitionings.intern` no longer
takes its return-as-is path for an inner join with one shuffled side: every
`outputPartitioning` call (a `def`, re-evaluated by `EnsureRequirements`,
`ValidateRequirements` and every ancestor join) rebuilds one member via
`keyed.copy(layout = canonicalLayout)`. When the marked side happens to be the
first member and the other side is a nested `PartitioningCollection`, the whole
subtree is rebuilt and `checkKeyedPartitioningInvariant` re-runs, so O(subtree)
where the base was O(1). In the base, `k.copy(mayContainUnknownPartitionKeys =
false)` kept the `partitionKeys` reference and both flags agreed, so `intern`
returned the member untouched. The `fromPartitionings` comment still promises
"returned as it is ... O(1) per level".
One cheap way back: when exactly one side is unmarked, reuse that side's
layout object as `cleared` if `k.layout.copy(mayContainUnknownPartitionKeys =
false) == unmarkedLayout` (O(1), since the case-class `==` short-circuits on
the shared `partitionKeys` reference), and fall back to the fresh copy
otherwise. Not a blocker; alternatively just qualify the comment in
`fromPartitionings`.
##########
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:
Nit: `sortKeys = true` has exactly one caller (`DataSourceV2ScanExecBase`),
and `KeyedPartitioning.groupedKeyRowOrdering(dataTypes)` already returns this
same factory ordering
(`getInternalRowComparableWrapperFactory(dataTypes).ordering`, cache-backed)
and is what `GroupPartitionsExec.groupAndSortByKeys` and `EnsureRequirements`
already sort with. The scan could do
`.sorted(KeyedPartitioning.groupedKeyRowOrdering(exprs.map(_.dataType)))` and
call the unchanged two-arg `apply`; both resolve to
`orderingCache.get(comparableTypes(...))`, so the "one derivation" property
holds by construction either way, and the sort contract is visibly shared with
`GroupPartitionsExec` by helper name rather than by a boolean 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:
Nit: `dataTypes`, `isGrouped` and `isCollapsed` are read only through
`layout` (every other `grouping.` access in the file is `partitions`,
`keysRewritten` or the two counters), and `outputPartitioning` re-reads the
marker separately and re-applies it with `if (marked) grouping.layout.copy(...)
else grouping.layout`, plus a comment explaining the split. `childKp` is in
scope where the grouping is built, so the layout could be built once there with
`mayContainUnknownPartitionKeys = childKp.mayContainUnknownPartitionKeys`, and
`PartitionGrouping` could hold `layout: KeyLayout` instead of the three fields.
The `UnknownPartitioning` give-up branch never reports the layout, so a pre-set
marker is unobservable there, and both branches already force `grouping`. That
removes the `marked` if/else and keeps one place to touch when `KeyLayout`
grows a field.
##########
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:
Nit: the inner `if k.mayContainUnknownPartitionKeys` guard is dead here,
since the outer `keyedMarkerOf(partitioning).contains(true)` already read the
shared marker and the constructor's `layout eq` invariant means every keyed
member carries it, and the `var cleared = null` first-visit dance is a lazy way
of computing the representative's layout once. It could match the style
`GroupPartitionsExec.outputPartitioning` already uses, computing the layout
outside the transform:
```scala
val cleared = partitioning
.collectFirst { case k: KeyedPartitioning => k.layout }.get
.copy(mayContainUnknownPartitionKeys = false)
partitioning.transform { case k: KeyedPartitioning => k.copy(layout =
cleared) }
```
(or widen `representativeOf` from `private[physical]` to `private[sql]`,
like `keyedMarkerOf` next to it).
##########
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:
Nit: the "same key space" predicate (`keyDataTypes ==` alongside
`partitionKeys ==`), and the "two empty key lists compare equal whatever they
describe" rationale, now live in two places: the requires here (1315 and 1324,
which also mix the two names `keyDataTypes` and `dataTypes` for the same field)
and `KeyedShuffleSpec.isCompatibleWith`, where the last round grew the comment
further. Since the empty-key argument is what motivated `KeyLayout` in the
first place, it might belong on it, e.g. `KeyLayout.describesSameKeys(other:
KeyLayout): Boolean = dataTypes == other.dataTypes && partitionKeys ==
other.partitionKeys`, with the rationale written once. `intern` would
`require(representative.layout.describesSameKeys(canonicalLayout))` plus the
separate `isGrouped` require (which is deliberately not part of
`isCompatibleWith`), and `isCompatibleWith` would call the same method, so the
next site that compares keys cannot forget the type clause.
##########
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:
Nit: `copy(expressions = ...).withLayout(...)` builds an intermediate
`KeyedPartitioning` (a `TreeNode`, so it captures `origin` and allocates the
lazy holders) that is discarded immediately. `partitioning.copy(expressions =
newExpressions, layout =
partitioning.layout.copy(mayContainUnknownPartitionKeys = true))` does it in
one node. Planning-time only, once per shuffle candidate, so purely a nit; feel
free to keep the `withLayout` idiom if you prefer the readability.
--
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]
