peter-toth commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4359560138
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,7 +716,129 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition information down to the scan nodes
+ /**
+ * 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)],
+ 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)
+ }
+
+ /**
+ * 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 =>
+ Some(projectAndDedupeCommonPartitionValues(
+ values, layout, positions, existing.applyPartialClustering))
+ case (None, values, _, _) => values
+ }
+ }
+
+ /**
+ * Reducers for a scan after a push. The reducers an inner join pushed index
the scan's reported
+ * layout, so they are first projected onto `newPositions`. A reducer this
join pushes for a
+ * position applies to values the inner reducer already produced, so the two
are chained.
+ */
+ private def composeReducers(
+ existing: Option[Seq[Option[Reducer[_, _]]]],
+ newPositions: Option[Seq[Int]],
+ pushed: Option[Seq[Option[Reducer[_, _]]]]):
Option[Seq[Option[Reducer[_, _]]]] = {
+ val projectedExisting = (existing, newPositions) match {
+ case (Some(rs), Some(positions)) => Some(positions.map(rs))
+ case _ => existing
+ }
+ (projectedExisting, pushed) match {
+ case (Some(inner), Some(outer)) =>
+ assert(inner.length == outer.length,
+ "pushed reducers must index the same layout as the composed join-key
positions")
+ Some(inner.zip(outer).map {
+ case (Some(i), Some(o)) => Some(ChainedReducer(i, o))
+ case (i, o) => o.orElse(i)
+ })
+ case (inner, outer) => outer.orElse(inner)
+ }
+ }
+
+ /**
+ * Whether `spec` uses a partition expression of a scan below `plan` whose
splits an inner join
+ * already reduced at that position. Such a scan reports its own partition
transform there but
+ * values of the reduced transform, so a spec built on that position
describes a layout the scan
+ * does not have. Positions without a reducer still hold the scan's own
values and are sound.
+ */
+ private def isBuiltFromReducedScan(plan: SparkPlan, spec:
KeyGroupedShuffleSpec): Boolean = {
+ val specExpressions = spec.partitioning.expressions
+ // Only a nested join can hand a reduced scan's layout to a later join.
Looking below a join
+ // also keeps a re-run of this rule over an already planned join from
matching its own scans.
+ plan.exists {
+ case join @ (_: SortMergeJoinExec | _: ShuffledHashJoinExec) =>
join.exists {
+ case scan: BatchScanExec =>
+ (reportedLayout(scan.spjParams), scan.spjParams.reducers) match {
+ case (Some(layout), Some(reducers)) =>
+ layout.zip(reducers).exists { case (e, r) =>
+ r.isDefined && specExpressions.exists(_.semanticEquals(e))
Review Comment:
**Finding 19.** This matches the spec's partition expressions to the scan's
own with `semanticEquals`, so an alias between the two joins hides the
reduction. `ProjectExec` rewrites the nested join's `KeyGroupedPartitioning`
through the alias, to `bucket(12, id2)`, which is not `semanticEquals` to the
scan's `bucket(12, id)`.
With `l12`, `l4` and `l3`, bucketed 12, 4 and 3 over `BIGINT` ids 0 to 23,
and `allowCompatibleTransforms` on:
```sql
SELECT /*+ MERGE(x, l3) */ x.id2, x.id4, l3.id
FROM (SELECT /*+ MERGE(l12, l4) */ l12.id AS id2, l4.id AS id4
FROM l12 JOIN l4 ON l12.id = l4.id) x
JOIN l3 ON x.id2 = l3.id
```
This returns 8 rows instead of 24, AQE on and off, with no shuffle, the same
as on `branch-4.1`. Without the alias it now returns all 24 rows with two
shuffles.
Only a transform position can hold a reducer, so comparing the transform
rather than its column closes it. With this, the query returns 24 rows with two
shuffles, and the suite passes:
```suggestion
r.isDefined && specExpressions.exists {
case s: TransformExpression => e match {
case t: TransformExpression => s.isSameFunction(t)
case _ => false
}
case s => s.semanticEquals(e)
}
```
It can also shuffle a side that is fine, when another column of a scan below
carries the same transform. This is the stale-transform problem
[SPARK-59121](https://issues.apache.org/jira/browse/SPARK-59121) fixes on 4.2.1
and later: a second join derives its reducer from the transform a reduced scan
still reports.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -229,7 +239,8 @@ case class EnsureRequirements(
// partitioned side's BatchScanExec is grouped by join keys to
match,
// and we do that by pushing down the join keys
case Some(KeyGroupedShuffleSpec(_, _, Some(joinKeyPositions))) =>
- populateJoinKeyPositions(child, Some(joinKeyPositions))
+ populateSpjParams(
Review Comment:
**Finding 20.** The reduced-scan check covers the best spec (L148-159) and
`checkKeyGroupCompatible` (L516), but not the children `matchedIndexes` keeps
(L212-214). When `checkKeyGroupCompatible` refuses the reduced member, the best
spec can come from the other side. A nested join whose reduced member
`isCompatibleWith` that spec is then kept unshuffled, and this push only
narrows its positions.
With `nb4` and `nb12` bucketed 4 and 12 over `BIGINT` ids 0 to 23, and `nc6`
bucketed 6 and holding ids 0, 1, 2, 3 and 6, with `shuffle.enabled` and
`allowCompatibleTransforms` on:
```sql
SELECT /*+ MERGE(nb4, nb12, nc6) */ nb4.id, nb12.id, nc6.id
FROM nb4 JOIN nb12 ON nb4.id = nb12.id
JOIN nc6 ON nb12.id = nc6.id
```
This returns 4 rows without id 6, AQE on and off, with no shuffle, the same
as on `branch-4.1`. Id 6 sits in partition 2 of the reduced side but in
partition 0 of `nc6`.
Leaving such a child out of `matchedIndexes` shuffles it onto `nc6`'s
layout. With this, the query returns 5 rows with one shuffle, and the suite
passes:
```scala
lazy val matchedIndexes = bestSpecOpt.toSeq.flatMap { best =>
childrenIndexes.filter(i => best.isCompatibleWith(specs(i)) &&
!flattenSpec(specs(i)).exists {
case k: KeyGroupedShuffleSpec =>
isBuiltFromReducedScan(children(i), k)
case _ => false
})
}
```
This is the same
[SPARK-59121](https://issues.apache.org/jira/browse/SPARK-59121) problem: a
reduced member is found compatible on the transform it still reports.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,7 +716,129 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition information down to the scan nodes
+ /**
+ * 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)],
+ 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)
+ }
+
+ /**
+ * 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 =>
+ Some(projectAndDedupeCommonPartitionValues(
+ values, layout, positions, existing.applyPartialClustering))
+ case (None, values, _, _) => values
+ }
+ }
+
+ /**
+ * Reducers for a scan after a push. The reducers an inner join pushed index
the scan's reported
+ * layout, so they are first projected onto `newPositions`. A reducer this
join pushes for a
+ * position applies to values the inner reducer already produced, so the two
are chained.
+ */
+ private def composeReducers(
+ existing: Option[Seq[Option[Reducer[_, _]]]],
+ newPositions: Option[Seq[Int]],
+ pushed: Option[Seq[Option[Reducer[_, _]]]]):
Option[Seq[Option[Reducer[_, _]]]] = {
+ val projectedExisting = (existing, newPositions) match {
+ case (Some(rs), Some(positions)) => Some(positions.map(rs))
+ case _ => existing
+ }
+ (projectedExisting, pushed) match {
+ case (Some(inner), Some(outer)) =>
+ assert(inner.length == outer.length,
+ "pushed reducers must index the same layout as the composed join-key
positions")
+ Some(inner.zip(outer).map {
+ case (Some(i), Some(o)) => Some(ChainedReducer(i, o))
+ case (i, o) => o.orElse(i)
+ })
+ case (inner, outer) => outer.orElse(inner)
+ }
+ }
+
+ /**
+ * Whether `spec` uses a partition expression of a scan below `plan` whose
splits an inner join
+ * already reduced at that position. Such a scan reports its own partition
transform there but
+ * values of the reduced transform, so a spec built on that position
describes a layout the scan
+ * does not have. Positions without a reducer still hold the scan's own
values and are sound.
+ */
+ private def isBuiltFromReducedScan(plan: SparkPlan, spec:
KeyGroupedShuffleSpec): Boolean = {
+ val specExpressions = spec.partitioning.expressions
+ // Only a nested join can hand a reduced scan's layout to a later join.
Looking below a join
+ // also keeps a re-run of this rule over an already planned join from
matching its own scans.
+ plan.exists {
+ case join @ (_: SortMergeJoinExec | _: ShuffledHashJoinExec) =>
join.exists {
+ case scan: BatchScanExec =>
+ (reportedLayout(scan.spjParams), scan.spjParams.reducers) match {
+ case (Some(layout), Some(reducers)) =>
+ layout.zip(reducers).exists { case (e, r) =>
+ r.isDefined && specExpressions.exists(_.semanticEquals(e))
+ }
+ case _ => false
+ }
+ case _ => false
+ }
+ case _ => false
+ }
+ }
+
+ // Populate the common partition information down to the scan nodes,
overwriting SPJ params.
Review Comment:
**Finding 18.** `ensureOrdering` (L361) is the only caller left, and the
reset is wrong there whenever the plan below already carries a join's SPJ
params. A LEFT OUTER join reports its left side's `KeyGroupedPartitioning`, so
an `ORDER BY` on the join key reaches this path. The reset then drops the key
positions and reducers the join pushed, and the scans are regrouped on the full
partition keys with values of the join keys.
What I ran, AQE on and off, with `bucketing.sorting.enabled`,
`allowJoinKeysSubsetOfPartitionKeys` and
`requireAllClusterKeysForCoPartition=false`:
```sql
-- oa3 PARTITIONED BY (a, b): (1, 10), (2, 20), (3, 10), (4, 30)
-- ob3 PARTITIONED BY (b): (0, 10), (0, 20)
SELECT /*+ MERGE(oa3, ob3) */ oa3.b, oa3.a, ob3.a
FROM oa3 LEFT JOIN ob3 ON oa3.b = ob3.b
ORDER BY oa3.b
```
- At this head and on `branch-4.1`: no rows with AQE off, and
`ArrayIndexOutOfBoundsException: Index 1 out of bounds for length 1` with AQE
on.
- With the change below: the 4 expected rows in both modes, and
`KeyGroupedPartitioningSuite` still passes.
Finding 8 asked to keep this reset, which was the wrong call. Pushing the
sorted values with `populateSpjParams` keeps what the join pushed, and
`populateCommonPartitionInfo` then has no caller left:
```scala
val sorted = partitionValues.sorted(partitionOrdering)
populateSpjParams(plan, commonPartitionValues = Some(sorted.map((_,
1))),
joinKeyPositions = None)
```
No ticket covers this one. If you'd rather not fix it here, please file a
ticket for it.
--
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]
