peter-toth commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4402322070
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -502,6 +537,9 @@ case class EnsureRequirements(
if (specOpt.isEmpty) return None
specOpt.get
}
+ if (Seq(left, right).zip(specs).exists { case (p, s) =>
isBuiltFromReducedScan(p, s) }) {
Review Comment:
**Finding 32.** The description says `checkKeyGroupCompatible` returns
`None` for a spec on a scan that `hasReducedScanOutsideJoin` finds, but only
`candidateSpecs` checks it. When `OptimizeSkewedJoin` re-applies this rule, the
inner join of a nested SPJ finds its scans directly below it, still carrying
the reducers and values of the first pass. `isBuiltFromReducedScan` needs a
join below the child, so it lets them through. The join then derives the
reducer again from the transform the scan still reports, `reducePartValues`
applies it to values it already produced, and `composeReducers` chains the two
into `ChainedReducer(r, r)`. That is only right when `r(r(x)) == r(x)`.
Bucket's `x % n` has that property, and the skew test uses only bucket reducers.
What I ran: I registered `hours`, `days` and `years` functions in the test
catalog that implement `ReducibleFunction`. They reduce hours to days with
`floorDiv(h, 24)`, and days to years, which is the shape of the
`ReducibleFunction` javadoc's days/hours example. `qh`, `qd` and `qy` are
`PARTITIONED BY (hours|days|years(ts))` and hold the same 72 hourly rows from
2026-01-01. With the confs of `nested SPJ survives EnsureRequirements
re-running for a skewed join`:
```sql
SELECT key1 AS k, count(*) AS c FROM skewData1 JOIN skewData2 ON key1 = key2
GROUP BY key1
UNION ALL
SELECT /*+ MERGE(qh, qd, qy) */ qh.id AS k, 1L AS c
FROM qh JOIN qd ON qh.ts = qd.ts JOIN qy ON qd.ts = qy.ts
```
- With `forceOptimizeSkewedJoin=true`, the second branch returns 0 of 72
rows.
- With `false`, or without the skewed branch (AQE on and off), it returns
all 72.
- `branch-4.1` returns 0 rows in all of these.
Checking it here too fixes the rows:
```suggestion
if (Seq(left, right).zip(specs).exists { case (p, s) =>
isBuiltFromReducedScan(p, s) } ||
Seq(left, right).exists(hasReducedScanOutsideJoin)) {
```
With it, the query returns all 72 rows, and both suites pass. The re-run's
inner join then shuffles too, so the skew test's query plans 6 shuffles instead
of 5. Please update that count under the user-facing changes. A test with a
reducer that is not idempotent would pin this, since bucket reducers cannot.
These are the functions I registered with `catalog.createFunction` before
creating the tables:
<details><summary>Test functions</summary>
```scala
import java.time.{Instant, LocalDate, ZoneOffset}
import java.time.temporal.ChronoUnit
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.types._
// hours reduce onto days with floorDiv(h, 24), and days onto years. Neither
reducer is
// idempotent, unlike bucket's x % n.
object ProbeHoursFunction extends ScalarFunction[Long] with
ReducibleFunction[Long, Long] {
override def inputTypes(): Array[DataType] = Array(TimestampType)
override def resultType(): DataType = LongType
override def name(): String = "hours"
override def canonicalName(): String = name()
override def produceResult(input: InternalRow): Long =
ChronoUnit.HOURS.between(Instant.EPOCH,
DateTimeUtils.microsToInstant(input.getLong(0)))
override def reducer(other: ReducibleFunction[_, _]): Reducer[Long, Long] =
if (other eq ProbeDaysFunction) ProbeHoursToDays else null
}
object ProbeDaysFunction extends ScalarFunction[Long] with
ReducibleFunction[Long, Long] {
override def inputTypes(): Array[DataType] = Array(TimestampType)
override def resultType(): DataType = LongType
override def name(): String = "days"
override def canonicalName(): String = name()
override def produceResult(input: InternalRow): Long =
ChronoUnit.DAYS.between(Instant.EPOCH,
DateTimeUtils.microsToInstant(input.getLong(0)))
override def reducer(other: ReducibleFunction[_, _]): Reducer[Long, Long] =
if (other eq ProbeYearsFunction) ProbeDaysToYears else null
}
object ProbeYearsFunction extends ScalarFunction[Long] with
ReducibleFunction[Long, Long] {
override def inputTypes(): Array[DataType] = Array(TimestampType)
override def resultType(): DataType = LongType
override def name(): String = "years"
override def canonicalName(): String = name()
override def produceResult(input: InternalRow): Long = {
val d =
DateTimeUtils.microsToInstant(input.getLong(0)).atZone(ZoneOffset.UTC).toLocalDate
ChronoUnit.YEARS.between(LocalDate.ofEpochDay(0), d)
}
override def reducer(other: ReducibleFunction[_, _]): Reducer[Long, Long]
= null
}
case object ProbeHoursToDays extends Reducer[Long, Long] {
override def reduce(hours: Long): Long = Math.floorDiv(hours, 24L)
}
case object ProbeDaysToYears extends Reducer[Long, Long] {
override def reduce(days: Long): Long =
ChronoUnit.YEARS.between(LocalDate.ofEpochDay(0),
LocalDate.ofEpochDay(days))
}
class ProbeUnbound(fn: BoundFunction) extends UnboundFunction {
override def bind(inputType: StructType): BoundFunction = fn
override def description(): String = name()
override def name(): String = fn.name()
}
```
</details>
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +740,202 @@ 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 =>
+ 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 = {
+ // A `Project` between the joins can rename the column under the
transform, so map the spec's
+ // columns back through the aliases below `plan` before comparing them
with the scan's. A
+ // broadcast hash join also reports the streamed side's partitioning on
its build keys, so map
+ // those back to the streamed keys too.
+ val broadcastKeys = plan.flatMap {
+ case j: BroadcastHashJoinExec =>
+ val (buildKeys, streamedKeys) = j.buildSide match {
+ case BuildLeft => (j.leftKeys, j.rightKeys)
+ case BuildRight => (j.rightKeys, j.leftKeys)
+ }
+ buildKeys.zip(streamedKeys).collect { case (b: Attribute, s:
Attribute) => b -> s }
+ case _ => Nil
+ }
+ val aliases = AttributeMap(plan.flatMap(_.expressions).flatMap(_.collect {
+ case a @ Alias(child: Attribute, _) => a.toAttribute -> child
+ }) ++ broadcastKeys)
+ def unalias(e: Expression): Expression = e.transformUp {
+ case a: Attribute if aliases.contains(a) => unalias(aliases(a))
+ }
+ val specExpressions = spec.partitioning.expressions.map(unalias)
Review Comment:
**Finding 33.** The build-key map I suggested in [finding
30](https://github.com/apache/spark/pull/59599#discussion_r4389493024) fails
open. It takes only `(Attribute, Attribute)` key pairs, and `AttributeMap`
keeps one streamed key per build key. A spec expression the map cannot follow
is then trusted as if no inner join had reduced it. Three shapes still read the
reduced `l12`'s stale `bucket(12)` layout, and each returns 8 of 24 rows with
no shuffle, AQE on and off:
- A build key equated to two streamed keys. With `JOIN lx ON lx.id = l12.id
AND lx.id = l4.id`, the map keeps `lx.id -> l4.id`, which nothing reduced. With
the two equalities swapped, the query returns all 24 rows.
- A cast build key. With `lxi.id` and `l3i.id` as `INT` and the other ids
`BIGINT`, the build key of `l12.id = lxi.id` is `cast(lxi.id as bigint)`:
```sql
SELECT /*+ MERGE(l12, l4), BROADCAST(lxi), MERGE(l3i) */ l12.id, lxi.id,
l3i.id
FROM l12 JOIN l4 ON l12.id = l4.id
JOIN lxi ON l12.id = lxi.id
JOIN l3i ON lxi.id = l3i.id
```
- A build key that a broadcast inside the build side also uses. `ly.id` is
the build key of both broadcast joins. The inner one is visited last, so the
map keeps `ly.id -> lx.id` and loses `ly.id -> l12.id`:
```sql
SELECT /*+ MERGE(l12, l4), BROADCAST(b), MERGE(l3) */ l12.id, b.id, l3.id
FROM l12 JOIN l4 ON l12.id = l4.id
JOIN (SELECT /*+ BROADCAST(ly) */ lx.id AS xid, ly.id
FROM lx JOIN ly ON lx.id = ly.id) b ON l12.id = b.id
JOIN l3 ON b.id = l3.id
```
The tables are those of finding 30, plus unpartitioned `lxi` and `ly` and
`bucket(3)` `l3i`, all with ids 0 to 23. `branch-4.1` returns the same 8 rows
for all three.
Mapping more cases would leave open the next rename the map does not know.
Reading `AliasAwareOutputExpression`, an alias of a struct field is one.
Failing closed covers them all: trust a spec expression only when it maps to a
scan position that no inner join reduced, and fall back to a shuffle otherwise.
That keeps the exact match finding 22 asked for. Walking the positions first
also lets it return before it builds the maps when nothing below carries a
reducer:
```scala
private def isBuiltFromReducedScan(plan: SparkPlan, spec:
KeyGroupedShuffleSpec): Boolean = {
// Each position of the scans that decide the layout above `plan`, with
whether an inner join
// reduced it. An exchange re-partitions its input, so a scan below it
no longer decides that
// layout. Only a nested join can hand a reduced scan's layout to a
later join, which also keeps
// a re-run of this rule from matching the scans directly below an
already planned join.
def positions(p: SparkPlan, belowJoin: Boolean): Seq[(Expression,
Boolean)] = p match {
case _: ShuffleExchangeLike | _: BroadcastExchangeLike => Nil
case scan: BatchScanExec =>
reportedLayout(scan.spjParams).toSeq.flatMap { layout =>
val reducers =
scan.spjParams.reducers.getOrElse(Seq.fill(layout.length)(None))
layout.zip(reducers).map { case (e, r) => (e, belowJoin &&
r.isDefined) }
}
case join @ (_: SortMergeJoinExec | _: ShuffledHashJoinExec) =>
join.children.flatMap(positions(_, belowJoin = true))
case other => other.children.flatMap(positions(_, belowJoin))
}
val scanPositions = positions(plan, belowJoin = false)
if (!scanPositions.exists(_._2)) return false
// A `Project` between the joins can rename the column under the
transform, so map the spec's
// columns back through the aliases below `plan` before comparing them
with the scan's. A
// broadcast hash join also reports the streamed side's partitioning on
its build keys, so map
// those back to the streamed keys too.
val broadcastKeys = plan.flatMap {
case j: BroadcastHashJoinExec =>
val (buildKeys, streamedKeys) = j.buildSide match {
case BuildLeft => (j.leftKeys, j.rightKeys)
case BuildRight => (j.rightKeys, j.leftKeys)
}
buildKeys.zip(streamedKeys).collect { case (b: Attribute, s:
Attribute) => b -> s }
case _ => Nil
}
val aliases = AttributeMap(plan.flatMap(_.expressions).flatMap(_.collect
{
case a @ Alias(child: Attribute, _) => a.toAttribute -> child
}) ++ broadcastKeys)
def unalias(e: Expression): Expression = e.transformUp {
case a: Attribute if aliases.contains(a) => unalias(aliases(a))
}
// Trust a spec expression only when it maps to a position no inner join
reduced. One the
// mapping above cannot follow, e.g. a cast build key, falls back to a
shuffle.
spec.partitioning.expressions.map(unalias).exists { e =>
!scanPositions.exists { case (p, reduced) => !reduced &&
p.semanticEquals(e) }
}
}
```
With it, all three queries return all 24 rows with two shuffles, and both
suites pass, including `keeps SPJ on an unreduced column with the same
transform`. If you'd rather leave this to SPARK-59121, I'll add these shapes
there.
--
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]
