peter-toth commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4412684753
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -502,6 +537,10 @@ case class EnsureRequirements(
if (specOpt.isEmpty) return None
specOpt.get
}
+ if (Seq(left, right).zip(specs).exists { case (p, s) =>
isBuiltFromReducedScan(p, s) } ||
+ Seq(left, right).exists(hasReducedScanOutsideJoin)) {
+ return None
+ }
Review Comment:
**Finding 34.** The `hasReducedScanOutsideJoin` check I suggested in finding
32 also catches a single SPJ with compatible transforms. When
`OptimizeSkewedJoin` re-applies this rule, `b4 JOIN b12` finds its own scans
directly below it. `b12`'s scan still carries the reducer the first run pushed,
so the join returns `None` here. With `v2.bucketing.shuffle.enabled` on, which
both skew tests set, the fallback still finds the two specs compatible and
keeps SPJ. With it off, which is the default, no spec can create a
partitioning, so both sides get a hash shuffle.
What I ran: `<skewed query> UNION ALL <SPJ>`, with the skewed query and the
confs of the two skew tests, and `shuffle.enabled` set as in the table. `b4`,
`b12` and `b8` are `bucket(4|12|8, id)` with ids 0 to 23, and `qh`, `qd`, `qy`
are the tables of `re-running EnsureRequirements does not apply a reducer
twice`. Without the skewed branch, `b4 JOIN b12` and `qh JOIN qd` plan no
shuffle.
| SPJ next to the skewed join | `shuffle.enabled` | `branch-4.1` | this head
| with the change below |
|---|---|---|---|---|
| `b4 JOIN b12`, `b4 JOIN b8`, or `b8` through an alias joined to `b4` |
false | 3 | 5 | 3 |
| the same | true | 3 | 3 | 3 |
| `qh JOIN qd`, `qd JOIN qy` | false | 3 | 5 | 3 |
| the same | true | 3 | 3 | 3 |
| the query of `nested SPJ survives EnsureRequirements re-running for a
skewed join` | both | 3 | 6 | 5 |
All return the right rows. That includes `branch-4.1` with hours and days,
whose reducer is not idempotent, so for a single SPJ the check fixes nothing.
The first run already pushed this join's pairing into both sides, so keeping
it is enough:
```suggestion
if (Seq(left, right).zip(specs).exists { case (p, s) =>
isBuiltFromReducedScan(p, s) }) {
return None
}
// A scan that carries reducers with no join above it in a child was
planned by an earlier
// run of this rule, which already pushed this join's pairing into both
sides. Keep that
// pairing rather than derive the reducers again.
if (Seq(left, right).exists(hasReducedScanOutsideJoin)) {
return Some(Seq(left, right))
}
```
With it, both suites pass, and the reducer test still returns all 72 rows.
The nested skew query plans 5 shuffles instead of 6, since its outer join still
refuses the reducer it pushed into the inner scans the first time. Please
update that count and the `checkKeyGroupCompatible` sentence in the
description. A test for a single `b4 JOIN b12` next to the skewed join, with
`shuffle.enabled` off and the shuffle count asserted, would pin this. With
`shuffle.enabled` on, as both skew tests set it now, the fallback hides it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +741,204 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition information down to the scan nodes
- private def populateCommonPartitionInfo(
- plan: SparkPlan,
+ /**
+ * Compose join-key positions from an outer SPJ with positions already
pushed by an inner SPJ.
+ * Outer positions index the layout the outer join sees; inner positions
index the scan's raw
+ * partition expressions, so they are composed rather than replaced.
+ */
+ private def composeJoinKeyPositions(
+ existing: Option[Seq[Int]],
+ newPositions: Option[Seq[Int]]): Option[Seq[Int]] = {
+ (existing, newPositions) match {
+ case (Some(old), Some(newPos)) => Some(newPos.map(old))
+ case (_, None) => existing
+ case (None, Some(newPos)) => Some(newPos)
+ }
+ }
+
+ /**
+ * Project common partition values onto `projectionPositions` of the layout
`expressions`
+ * describes and deduplicate. Values that collapse to the same key are read
as one group, so
+ * they become one partition, unless the scan is partially clustered, where
each split of the
+ * group still gets its own partition and the split counts add up.
+ *
+ * The keys keep the order in which they first appear, as
`KeyGroupedPartitioning.apply` does.
+ * A one-side shuffle routes the other child into that order, so partition i
of both sides must
+ * hold the same key.
+ */
+ private def projectAndDedupeCommonPartitionValues(
values: Seq[(InternalRow, Int)],
- joinKeyPositions: Option[Seq[Int]],
- reducers: Option[Seq[Option[Reducer[_, _]]]],
- applyPartialClustering: Boolean,
- replicatePartitions: Boolean): SparkPlan = plan match {
- case scan: BatchScanExec =>
- scan.copy(
- spjParams = scan.spjParams.copy(
- commonPartitionValues = Some(values),
- joinKeyPositions = joinKeyPositions,
- reducers = reducers,
- applyPartialClustering = applyPartialClustering,
- replicatePartitions = replicatePartitions
- )
- )
- case node =>
- node.mapChildren(child => populateCommonPartitionInfo(
- child, values, joinKeyPositions, reducers, applyPartialClustering,
replicatePartitions))
+ expressions: Seq[Expression],
+ projectionPositions: Seq[Int],
+ partiallyClustered: Boolean): Seq[(InternalRow, Int)] = {
+ val projectedExpressions = projectionPositions.map(expressions(_))
+ val projected = values.map { case (partValue, numSplits) =>
+ val row = KeyGroupedPartitioning.project(expressions,
projectionPositions, partValue)
+ (InternalRowComparableWrapper(row, projectedExpressions), numSplits)
+ }
+ val splitsPerKey = projected.groupMapReduce(_._1)(_._2)(_ + _)
+ projected.map(_._1).distinct.map { wrapper =>
+ (wrapper.row, if (partiallyClustered) splitsPerKey(wrapper) else 1)
+ }
Review Comment:
**Finding 35.** `partiallyClustered` is never true here. The projection runs
only for a push that brings no values, which is the positions-only push into a
compatible child at L251-252, and only into a scan an inner join already pushed
values to. A partially clustered scan reports `isPartiallyClustered`, so its
partitioning satisfies no `ClusteredDistribution` (`partitioning.scala:413`). A
later join therefore puts a shuffle above it, and the push stops there. So the
split counts can go:
```suggestion
projected.map(_._1).distinct.map(wrapper => (wrapper.row, 1))
```
The parameter, the doc comment's "unless the scan is partially clustered"
clause, and the description's "Split counts are added up only when the scan is
partially clustered" can go with them. This one is from reading, and both
suites pass with it and the change in 36.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +741,204 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition information down to the scan nodes
- private def populateCommonPartitionInfo(
- plan: SparkPlan,
+ /**
+ * Compose join-key positions from an outer SPJ with positions already
pushed by an inner SPJ.
+ * Outer positions index the layout the outer join sees; inner positions
index the scan's raw
+ * partition expressions, so they are composed rather than replaced.
+ */
+ private def composeJoinKeyPositions(
+ existing: Option[Seq[Int]],
+ newPositions: Option[Seq[Int]]): Option[Seq[Int]] = {
+ (existing, newPositions) match {
+ case (Some(old), Some(newPos)) => Some(newPos.map(old))
+ case (_, None) => existing
+ case (None, Some(newPos)) => Some(newPos)
+ }
+ }
+
+ /**
+ * Project common partition values onto `projectionPositions` of the layout
`expressions`
+ * describes and deduplicate. Values that collapse to the same key are read
as one group, so
+ * they become one partition, unless the scan is partially clustered, where
each split of the
+ * group still gets its own partition and the split counts add up.
+ *
+ * The keys keep the order in which they first appear, as
`KeyGroupedPartitioning.apply` does.
+ * A one-side shuffle routes the other child into that order, so partition i
of both sides must
+ * hold the same key.
+ */
+ private def projectAndDedupeCommonPartitionValues(
values: Seq[(InternalRow, Int)],
- joinKeyPositions: Option[Seq[Int]],
- reducers: Option[Seq[Option[Reducer[_, _]]]],
- applyPartialClustering: Boolean,
- replicatePartitions: Boolean): SparkPlan = plan match {
- case scan: BatchScanExec =>
- scan.copy(
- spjParams = scan.spjParams.copy(
- commonPartitionValues = Some(values),
- joinKeyPositions = joinKeyPositions,
- reducers = reducers,
- applyPartialClustering = applyPartialClustering,
- replicatePartitions = replicatePartitions
- )
- )
- case node =>
- node.mapChildren(child => populateCommonPartitionInfo(
- child, values, joinKeyPositions, reducers, applyPartialClustering,
replicatePartitions))
+ expressions: Seq[Expression],
+ projectionPositions: Seq[Int],
+ partiallyClustered: Boolean): Seq[(InternalRow, Int)] = {
+ val projectedExpressions = projectionPositions.map(expressions(_))
+ val projected = values.map { case (partValue, numSplits) =>
+ val row = KeyGroupedPartitioning.project(expressions,
projectionPositions, partValue)
+ (InternalRowComparableWrapper(row, projectedExpressions), numSplits)
+ }
+ val splitsPerKey = projected.groupMapReduce(_._1)(_._2)(_ + _)
+ projected.map(_._1).distinct.map { wrapper =>
+ (wrapper.row, if (partiallyClustered) splitsPerKey(wrapper) else 1)
+ }
}
+ /**
+ * The partition expressions a scan currently reports, i.e. its own
partition expressions
+ * projected onto the join-key positions pushed so far.
+ */
+ private def reportedLayout(params: StoragePartitionJoinParams):
Option[Seq[Expression]] =
+ params.keyGroupedPartitioning.map { expressions =>
+ params.joinKeyPositions.map(_.map(expressions)).getOrElse(expressions)
+ }
- private def populateJoinKeyPositions(
+ /**
+ * Common partition values for a scan after a push. Values pushed by this
join are the merge with
+ * the other side and already sit in the layout `newPositions` select, so
they win. Otherwise the
+ * values an inner join pushed sit in the scan's reported layout, and are
projected onto
+ * `newPositions`, which index that layout.
+ */
+ private def finalizeCommonPartitionValues(
+ existing: StoragePartitionJoinParams,
+ pushedValues: Option[Seq[(InternalRow, Int)]],
+ newPositions: Option[Seq[Int]]): Option[Seq[(InternalRow, Int)]] = {
+ (pushedValues, existing.commonPartitionValues, newPositions,
reportedLayout(existing)) match {
+ case (Some(values), _, _, _) => Some(values)
+ case (None, Some(values), Some(positions), Some(layout))
+ if values.nonEmpty && values.head._1.numFields == layout.length &&
+ positions != layout.indices =>
Review Comment:
**Finding 36.** When the width check fails, this falls through to `case
(None, values, _, _) => values` and keeps values in the old layout, while
`populateSpjParams` still composes the new positions. Every push puts its
values in the layout the scan reports, so the check always holds. If it ever
did not, the scan would group on keys it does not report. An assert, like the
one in `composeReducers`, would say so instead:
```suggestion
case (None, Some(values), Some(positions), Some(layout)) if positions
!= layout.indices =>
assert(values.forall(_._1.numFields == layout.length),
"common partition values must sit in the scan's reported layout")
```
Both suites pass with it and the change in 35.
--
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]
