peter-toth commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4424916439
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -345,10 +365,25 @@ case class EnsureRequirements(
},
Nil)
// Sort 'commonPartitionValues' and use this mechanism to ensure
BatchScan's
- // output partitions are ordered
+ // output partitions are ordered. A join below may already have pushed
key positions and
+ // reducers that these values are expressed in, so keep them rather
than reset them.
val sorted = partitionValues.sorted(partitionOrdering)
- populateCommonPartitionInfo(plan, sorted.map((_, 1)),
- None, None, applyPartialClustering = false, replicatePartitions =
false)
+ // A one-side shuffle below routed its other child into the current
order of the values.
+ // Sorting only the scans would leave that child behind, so fall back
to a range shuffle.
+ val hasKeyGroupedShuffle = plan.exists {
Review Comment:
**Finding 37.** `populateSpjParams` stops at a broadcast, but this walk goes
on into the build side. So a one-side shuffle inside a broadcast build side
also triggers the range shuffle, though the sorted push never reaches it.
`sa` and `ba` are `PARTITIONED BY (dept)` and hold depts 1 to 3. `bb` has
`dept` and `v`, is not partitioned, and holds `(1, 'v1')`, `(2, 'v2')` and `(3,
'v3')`. With `v2.bucketing.shuffle.enabled` and `v2.bucketing.sorting.enabled`:
```sql
SELECT /*+ BROADCAST(b) */ sa.dept, b.v
FROM sa LEFT JOIN (
SELECT /*+ MERGE(ba, bb) */ ba.dept, bb.v
FROM ba JOIN bb ON ba.dept = bb.dept) b
ON sa.dept = b.dept
ORDER BY sa.dept DESC
```
- With AQE off, this head plans a range shuffle in addition to `bb`'s
one-side shuffle, and returns the right rows. With AQE on, it plans only `bb`'s
shuffle.
- `branch-4.1` returns `[3,null]`, `[2,v2]`, `[1,null]`, with AQE on and
off, since its push reaches `ba`.
The `DESC` makes the sort reorder the values. Stopping the walk where the
push stops avoids the extra shuffle:
```scala
def hasKeyGroupedShuffle(p: SparkPlan): Boolean = p match {
case s @ (_: ShuffleExchangeLike | _: ShuffleQueryStageExec) =>
s.outputPartitioning.isInstanceOf[KeyGroupedPartitioning]
case _: BroadcastExchangeLike => false
case other => other.children.exists(hasKeyGroupedShuffle)
}
```
A shuffle that is not key-grouped also ends the walk there, since the push
does not go below it either. With this, the query plans only `bb`'s shuffle
with AQE off too, and returns the same rows.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +746,198 @@ 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.
+ *
+ * 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(
Review Comment:
**Finding 40.** Since the change for finding 35, this is exactly the
project-and-dedupe of `KeyGroupedPartitioning.apply`
(`partitioning.scala:485`), whose order the doc comment says it must match.
Calling it makes that hold by construction:
```scala
private def projectAndDedupeCommonPartitionValues(
values: Seq[(InternalRow, Int)],
expressions: Seq[Expression],
projectionPositions: Seq[Int]): Seq[(InternalRow, Int)] = {
KeyGroupedPartitioning(expressions, projectionPositions,
values.map(_._1), Nil,
isPartiallyClustered = false, mayContainUnknownPartitionKeys = false)
.partitionValues.map((_, 1))
}
```
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala:
##########
@@ -177,3 +177,53 @@ object TruncateFunction extends ScalarFunction[UTF8String]
{
str.substring(0, length)
}
}
+
+// Reducible `hours`, `days` and `years`, matching the in-memory table's
transforms of those
+// names. Hours reduce onto days with floorDiv(h, 24), and days onto years.
Unlike bucket's x % n,
+// neither reducer is idempotent, so applying one twice gives a wrong value.
+object ReducibleHoursFunction 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 ReducibleDaysFunction) HoursToDaysReducer else null
+}
+
+object ReducibleDaysFunction 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 ReducibleYearsFunction) DaysToYearsReducer else null
+}
+
+object ReducibleYearsFunction 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 =
+ YearsFunction.invoke(input.getLong(0))
+ override def reducer(other: ReducibleFunction[_, _]): Reducer[Long, Long] =
null
+}
+
+case object HoursToDaysReducer extends Reducer[Long, Long] {
+ override def reduce(hours: Long): Long = Math.floorDiv(hours, 24L)
Review Comment:
**Finding 41.** The in-memory table's `hours` and `days` transforms use
`ChronoUnit.HOURS.between` and `DAYS.between` (`InMemoryBaseTable.scala:229`
and `:236`), which truncate toward zero. `floorDiv` floors. So before 1970 the
reducer disagrees with `days`: `1969-12-31 23:00 UTC` has hours -1 and days 0,
while `floorDiv(-1, 24)` is -1. The tests only use 2026 timestamps, so nothing
fails today.
```suggestion
override def reduce(hours: Long): Long = hours / 24
```
The comment at line 182 would then say `h / 24`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +746,198 @@ 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.
+ *
+ * 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]): Seq[(InternalRow, Int)] = {
+ val projectedExpressions = projectionPositions.map(expressions(_))
+ val projected = values.map { case (partValue, _) =>
+ val row = KeyGroupedPartitioning.project(expressions,
projectionPositions, partValue)
+ InternalRowComparableWrapper(row, projectedExpressions)
+ }
+ projected.distinct.map(wrapper => (wrapper.row, 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 positions !=
layout.indices =>
+ assert(values.forall(_._1.numFields == layout.length),
+ "common partition values must sit in the scan's reported layout")
+ Some(projectAndDedupeCommonPartitionValues(values, layout, positions))
+ 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` may be built on 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. A spec expression that cannot be traced to one of those is treated
as unsound.
+ */
+ 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.
+ // `hasReducedScanOutsideJoin` covers those.
+ 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 }
Review Comment:
**Finding 39.** `AttributeMap` keeps one streamed key per build key, the
last one. So a build key equated to two streamed keys is traced to whichever
equality comes last. The expanded member can carry the build key in the other
key's position. When that position is reduced and the traced one is not, the
check trusts a stale layout.
`pa` is `PARTITIONED BY (bucket(12, a), bucket(12, b))` and `pb` is
`PARTITIONED BY (bucket(12, a), bucket(4, b))`, so the inner join reduces `pa`
at `b` only. `lx` is not partitioned, and `q12` is `PARTITIONED BY (bucket(12,
id))`. All hold BIGINT ids 0 to 23, with `a = b = id`. With
`allowCompatibleTransforms`, `allowJoinKeysSubsetOfPartitionKeys.enabled` and
`requireAllClusterKeysForCoPartition=false`:
```sql
SELECT /*+ MERGE(pa, pb), BROADCAST(lx), MERGE(q12) */ pa.a, lx.id, q12.id
FROM pa JOIN pb ON pa.a = pb.a AND pa.b = pb.b
JOIN lx ON lx.id = pa.b AND lx.id = pa.a
JOIN q12 ON lx.id = q12.id
```
- This head returns 8 of 24 rows with no shuffle, AQE on and off.
- With the two equalities swapped, the map keeps `lx.id -> pa.b`, and the
query falls back to 2 shuffles with all 24 rows.
- `branch-4.1` returns 8 of 24 rows in both orders.
The description lists a build key equated to two streamed keys among the
shapes this PR fixes. Its test joins `bucket(12)` and `bucket(4)` tables, so it
falls back in either order. Failing closed when a build key maps to more than
one streamed key covers both orders:
```scala
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
}.groupBy(_._1.exprId).values.collect {
// A build key equated to several streamed keys has no single position
to trace back to.
case pairs if pairs.map(_._2.exprId).distinct.size == 1 => pairs.head
}.toSeq
```
With it, the query returns all 24 rows with 2 shuffles in both orders, AQE
on and off. It costs a shuffle where the map happens to pick the right key,
which is the fail-closed trade the description already makes. Adding the query
above, in both orders, to `nested SPJ falls back when a broadcast build key
cannot be traced` would pin it.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -2980,6 +2980,1113 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase {
}
}
+ test("SPARK-60232: nested SPJ preserves inner join key positions") {
+ val cols = Array(
+ Column.create("store_id", IntegerType),
+ Column.create("dept_id", IntegerType),
+ Column.create("data", StringType))
+ // t4 is partitioned by (store_id, dept_id) and joins on dept_id only. The
inner join
+ // pushes joinKeyPositions [1] into the t4 scan. The outer join must not
overwrite that
+ // with [0] against the already projected dept_id, or t4 is grouped by
store_id and
+ // the three-way join returns no rows.
+ createTable("t4", cols, Array(identity("store_id"), identity("dept_id")))
+ sql("INSERT INTO testcat.ns.t4 VALUES (1, 2, 'a'), (2, 3, 'b'), (3, 1,
'c')")
+
+ createTable("t5", cols, Array(identity("dept_id")))
+ sql("INSERT INTO testcat.ns.t5 VALUES (0, 1, 'x1'), (0, 2, 'x2'), (0, 3,
'x3')")
+
+ createTable("t6", cols, Array(identity("dept_id")))
+ sql("INSERT INTO testcat.ns.t6 VALUES (0, 1, 'x1'), (0, 2, 'x2'), (0, 3,
'x3')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(t4, t5, t6) */
+ | t4.store_id, t4.dept_id, t4.data, t5.data, t6.data
+ |FROM testcat.ns.t4
+ |JOIN testcat.ns.t5 ON t4.dept_id = t5.dept_id
+ |JOIN testcat.ns.t6 ON t4.dept_id = t6.dept_id
+ |ORDER BY t4.store_id, t4.dept_id
+ |""".stripMargin)
+
+ val plan = df.queryExecution.executedPlan
+ assert(collectShuffles(plan).isEmpty, "SPJ should be triggered for
both joins")
+ val t4Scan = collectScans(plan).find(_.table.name().endsWith("t4")).get
+ assert(t4Scan.spjParams.joinKeyPositions.contains(Seq(1)),
+ "t4 scan should keep the inner join's key position [1], not [0]")
+ checkAnswer(df, Seq(
+ Row(1, 2, "a", "x2", "x2"),
+ Row(2, 3, "b", "x3", "x3"),
+ Row(3, 1, "c", "x1", "x1")))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ with mixed partition key types preserves key
positions") {
+ val cols = Array(
+ Column.create("id", StringType),
+ Column.create("date", StringType),
+ Column.create("data", StringType))
+ // Same nested join-key overwrite as above, but the leading partition key
has a different
+ // type than the bucket join key. Without the fix this shape throws
ClassCastException in
+ // connectors such as Iceberg; here the in-memory source only misaligns
values.
+ createTable("t1", cols,
+ Array(identity("date"), bucket(32, "id"), bucket(32, "data")))
+ sql("INSERT INTO testcat.ns.t1 VALUES " +
+ "('a', '2026-05-01', 'aa'), " +
+ "('a', '2026-05-02', 'cc'), " +
+ "('b', '2026-05-01', 'bb')")
+
+ createTable("t2", cols, Array(bucket(32, "id")))
+ sql("INSERT INTO testcat.ns.t2 VALUES " +
+ "('a', '2026-05-01', 'xx'), " +
+ "('b', '2026-05-01', 'yy')")
+
+ createTable("t3", cols, Array(bucket(32, "id")))
+ sql("INSERT INTO testcat.ns.t3 VALUES " +
+ "('a', '2026-05-01', 'xx'), " +
+ "('b', '2026-05-01', 'yy')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(t1, t2, t3) */
+ | t1.id, t1.date, t1.data, t2.data, t3.data
+ |FROM testcat.ns.t1
+ |JOIN testcat.ns.t2 ON t1.id = t2.id
+ |JOIN testcat.ns.t3 ON t1.id = t3.id
+ |ORDER BY t1.id, t1.date, t1.data
+ |""".stripMargin)
+
+ val plan = df.queryExecution.executedPlan
+ assert(collectShuffles(plan).isEmpty, "SPJ should be triggered for
both joins")
+ val t1Scan = collectScans(plan).find(_.table.name().endsWith("t1")).get
+ assert(t1Scan.spjParams.joinKeyPositions.contains(Seq(1)),
+ "t1 scan should keep the inner join's key position [1], not [0]")
+ checkAnswer(df, Seq(
+ Row("a", "2026-05-01", "aa", "xx", "xx"),
+ Row("a", "2026-05-02", "cc", "xx", "xx"),
+ Row("b", "2026-05-01", "bb", "yy", "yy")))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ preserves inner join reducers") {
+ val idCols = Array(Column.create("id", IntegerType))
+ createTable("b4", idCols, Array(bucket(4, "id")))
+ sql("INSERT INTO testcat.ns.b4 VALUES (4)")
+
+ createTable("b12", idCols, Array(bucket(12, "id")))
+ sql("INSERT INTO testcat.ns.b12 VALUES (4)")
+
+ createTable("b8", idCols, Array(bucket(8, "id")))
+ sql("INSERT INTO testcat.ns.b8 VALUES (1), (4)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(b4, b12, b8) */ b4.id, b12.id, b8.id
+ |FROM testcat.ns.b4
+ |JOIN testcat.ns.b12 ON b4.id = b12.id
+ |JOIN testcat.ns.b8 ON b4.id = b8.id
+ |""".stripMargin)
+
+ val plan = df.queryExecution.executedPlan
+ assert(collectShuffles(plan).isEmpty, "SPJ should be triggered for
both joins")
+ val scans = collectScans(plan)
+ def reducersOf(table: String): Option[Seq[Option[Reducer[_, _]]]] =
+
scans.find(_.table.name().endsWith(s".$table")).get.spjParams.reducers
+ assert(reducersOf("b4").isEmpty, "b4 is the coarsest side and needs no
reducer")
+ assert(reducersOf("b12").exists(_.length == 1) &&
reducersOf("b12").get.head.isDefined,
+ "b12 scan should keep the inner join's reducer onto bucket(4, id)")
+ assert(reducersOf("b8").exists(_.length == 1) &&
reducersOf("b8").get.head.isDefined,
+ "b8 scan should get the outer join's reducer onto bucket(4, id)")
+ checkAnswer(df, Seq(Row(4, 4, 4)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ preserves join key positions over common
partition values") {
+ val items2 = "items2"
+ val itemsPartitions = Array(identity("id"), identity("name"))
+
+ val inserts = Seq(
+ s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', 20.0, cast('2020-01-01' as timestamp))",
+ s"INSERT INTO testcat.ns.$items2 VALUES " +
+ "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " +
+ "(2, 'bb', 20.0, cast('2020-01-01' as timestamp))",
+ "INSERT INTO testcat.ns.purchases VALUES " +
+ "(1, 100.0, cast('2020-01-01' as timestamp)), " +
+ "(2, 200.0, cast('2020-01-01' as timestamp))")
+ val expected = Seq(Row(1, "aa", 10.0, 100.0), Row(2, "bb", 20.0, 200.0))
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ catalog.clearTables()
+ createTable(items, itemsColumns, itemsPartitions)
+ createTable(items2, itemsColumns, itemsPartitions)
+ createTable(purchases, purchasesColumns, Array.empty)
+ inserts.foreach(sql)
+
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(i, j, p) */ i.id, i.name, i.price, p.price
+ |FROM testcat.ns.$items i JOIN testcat.ns.$items2 j
+ |ON i.id = j.id AND i.name = j.name
+ |JOIN testcat.ns.$purchases p ON i.id = p.item_id
+ |ORDER BY i.id, i.name
+ |""".stripMargin)
+
+ val plan = df.queryExecution.executedPlan
+ assert(collectShuffles(plan).size == 1,
+ "only the unpartitioned purchases side should shuffle")
+ val itemsScan =
collectScans(plan).find(_.table.name().endsWith(items)).get
+ assert(itemsScan.spjParams.commonPartitionValues.nonEmpty,
+ "inner join should push common partition values into the items scan")
+ assert(itemsScan.spjParams.joinKeyPositions.contains(Seq(0)),
+ "outer join should project common partition values down to the id
join key")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ with multiple names per id over common
partition values") {
+ val items2 = "items2"
+ val itemsPartitions = Array(identity("id"), identity("name"))
+ val inserts = Seq(
+ s"INSERT INTO testcat.ns.$items VALUES " +
+ "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " +
+ "(1, 'bb', 20.0, cast('2020-01-02' as timestamp))",
+ s"INSERT INTO testcat.ns.$items2 VALUES " +
+ "(1, 'aa', 10.0, cast('2020-01-01' as timestamp)), " +
+ "(1, 'bb', 20.0, cast('2020-01-02' as timestamp))",
+ "INSERT INTO testcat.ns.purchases VALUES " +
+ "(1, 100.0, cast('2020-01-01' as timestamp))")
+ val expected = Seq(Row(1, "aa", 10.0, 100.0), Row(1, "bb", 20.0, 100.0))
+ Seq(true, false).foreach { aqe =>
+ catalog.clearTables()
+ createTable(items, itemsColumns, itemsPartitions)
+ createTable(items2, itemsColumns, itemsPartitions)
+ createTable(purchases, purchasesColumns, Array.empty)
+ inserts.foreach(sql)
+
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(i, j, p) */ i.id, i.name, i.price, p.price
+ |FROM testcat.ns.$items i JOIN testcat.ns.$items2 j
+ |ON i.id = j.id AND i.name = j.name
+ |JOIN testcat.ns.$purchases p ON i.id = p.item_id
+ |ORDER BY i.id, i.name
+ |""".stripMargin)
+
+ val plan = df.queryExecution.executedPlan
+ assert(collectShuffles(plan).size == 1,
+ "only the unpartitioned purchases side should shuffle")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+
+ private def createNestedSpjIdTables(): Unit = {
+ val cols = Array(
+ Column.create("store", IntegerType),
+ Column.create("id", IntegerType),
+ Column.create("data", StringType))
+ Seq("e1", "e2", "e3", "e5").foreach(t => createTable(t, cols,
Array(identity("id"))))
+ sql("INSERT INTO testcat.ns.e1 VALUES (0, 1, 'a'), (0, 2, 'b')")
+ sql("INSERT INTO testcat.ns.e2 VALUES (0, 1, 'a'), (0, 2, 'b')")
+ sql("INSERT INTO testcat.ns.e3 VALUES (0, 1, 'a'), (0, 2, 'b'), (0, 3,
'c')")
+ sql("INSERT INTO testcat.ns.e5 VALUES (0, 2, 'b'), (0, 3, 'c')")
+ createTable("e4", cols, Array(identity("store"), identity("id")))
+ sql("INSERT INTO testcat.ns.e4 VALUES (10, 1, 'a'), (20, 2, 'b'), (30, 3,
'c')")
+ }
+
+ test("SPARK-60232: nested SPJ pushes outer merged values into a nested
join") {
+ createNestedSpjIdTables()
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(e1, e2, e5) */ e1.id, e2.id, e5.id
+ |FROM testcat.ns.e1 JOIN testcat.ns.e2 ON e1.id = e2.id
+ |RIGHT JOIN testcat.ns.e5 ON e1.id = e5.id
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(2, 2, 2), Row(null, null, 3)))
+ }
+
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(e1, e2, e3) */ e1.id, e2.id, e3.id
+ |FROM testcat.ns.e1 JOIN testcat.ns.e2 ON e1.id = e2.id
+ |JOIN testcat.ns.e3 ON e1.id = e3.id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "SPJ should be triggered for both joins")
+ checkAnswer(df, Seq(Row(1, 1, 1), Row(2, 2, 2)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ with subset join keys on a two-key partitioned
side") {
+ createNestedSpjIdTables()
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(e1, e2, e4) */ e1.id, e2.id, e4.id, e4.store
+ |FROM testcat.ns.e1 JOIN testcat.ns.e2 ON e1.id = e2.id
+ |JOIN testcat.ns.e4 ON e1.id = e4.id
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(1, 1, 1, 10), Row(2, 2, 2, 20)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ on all partition keys, then on a subset") {
+ val cols = Array(
+ Column.create("store_id", IntegerType),
+ Column.create("dept_id", IntegerType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("store_id"), identity("dept_id")))
+ createTable("b", cols, Array(identity("store_id"), identity("dept_id")))
+ createTable("c", cols, Array(identity("dept_id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 2, 'a1'), (2, 2, 'a2'), (3, 1,
'a3')")
+ sql("INSERT INTO testcat.ns.b VALUES (1, 2, 'b1'), (2, 2, 'b2'), (3, 1,
'b3')")
+ sql("INSERT INTO testcat.ns.c VALUES (0, 1, 'c1'), (0, 2, 'c2'), (0, 3,
'c3')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val expected = Seq(
+ Row(1, 2, "a1", 1, 2, "b1", 0, 2, "c2"),
+ Row(2, 2, "a2", 2, 2, "b2", 0, 2, "c2"),
+ Row(3, 1, "a3", 3, 1, "b3", 0, 1, "c1"))
+ Seq("a.dept_id", "b.dept_id").foreach { outerKey =>
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(a, b, c) */ *
+ |FROM testcat.ns.a JOIN testcat.ns.b
+ |ON a.store_id = b.store_id AND a.dept_id = b.dept_id
+ |JOIN testcat.ns.c ON $outerKey = c.dept_id
+ |""".stripMargin)
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ dedupes common values that project to the same
join key") {
+ val cols = Array(
+ Column.create("store_id", IntegerType),
+ Column.create("dept_id", IntegerType),
+ Column.create("data", StringType))
+ createTable("a", cols, Array(identity("store_id"), identity("dept_id")))
+ createTable("b", cols, Array(identity("store_id"), identity("dept_id")))
+ createTable("c", cols, Array(identity("dept_id")))
+ sql("INSERT INTO testcat.ns.a VALUES (1, 2, 'a1'), (2, 2, 'a2')")
+ sql("INSERT INTO testcat.ns.b VALUES (1, 2, 'b1'), (2, 2, 'b2')")
+ sql("INSERT INTO testcat.ns.c VALUES (0, 1, 'c1'), (0, 2, 'c2'), (0, 3,
'c3')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(a, b, c) */ a.data, b.data, c.data
+ |FROM testcat.ns.a JOIN testcat.ns.b
+ |ON a.store_id = b.store_id AND a.dept_id = b.dept_id
+ |JOIN testcat.ns.c ON a.dept_id = c.dept_id
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row("a1", "b1", "c2"), Row("a2", "b2", "c2")))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ with several names per id and a partitioned
third side") {
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("name", StringType),
+ Column.create("price", DoubleType))
+ createTable("i", cols, Array(identity("id"), identity("name")))
+ createTable("j", cols, Array(identity("id"), identity("name")))
+ createTable("k", Array(Column.create("id", IntegerType),
Column.create("v", DoubleType)),
+ Array(identity("id")))
+ sql("INSERT INTO testcat.ns.i VALUES (1, 'aa', 1.0), (1, 'bb', 2.0)")
+ sql("INSERT INTO testcat.ns.j VALUES (1, 'aa', 1.0), (1, 'bb', 2.0)")
+ sql("INSERT INTO testcat.ns.k VALUES (1, 10.0), (2, 20.0)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ checkAnswer(
+ sql(
+ """
+ |SELECT /*+ MERGE(i, j, k) */ i.id, i.name, j.name, k.v
+ |FROM testcat.ns.i JOIN testcat.ns.j ON i.id = j.id AND i.name =
j.name
+ |JOIN testcat.ns.k ON i.id = k.id
+ |""".stripMargin),
+ Seq(Row(1, "aa", "aa", 10.0), Row(1, "bb", "bb", 10.0)))
+ checkAnswer(
+ sql(
+ """
+ |SELECT /*+ MERGE(i, j, k) */ i.id, i.name, k.v
+ |FROM testcat.ns.i LEFT ANTI JOIN testcat.ns.j
+ |ON i.id = j.id AND i.name = j.name
+ |JOIN testcat.ns.k ON i.id = k.id
+ |""".stripMargin),
+ Seq.empty)
+ }
+ }
+ }
+
+ test("SPARK-60232: four-way nested SPJ joining on a key dropped by an
earlier join") {
+ val cols = Array(
+ Column.create("id", IntegerType),
+ Column.create("name", StringType),
+ Column.create("price", DoubleType))
+ createTable("i", cols, Array(identity("id"), identity("name")))
+ createTable("j", cols, Array(identity("id"), identity("name")))
+ createTable("k", Array(Column.create("id", IntegerType),
Column.create("v", DoubleType)),
+ Array(identity("id")))
+ createTable("m", Array(Column.create("name", StringType),
Column.create("w", DoubleType)),
+ Array(identity("name")))
+ sql("INSERT INTO testcat.ns.i VALUES (1, 'aa', 1.0), (2, 'bb', 2.0)")
+ sql("INSERT INTO testcat.ns.j VALUES (1, 'aa', 1.0), (2, 'bb', 2.0)")
+ sql("INSERT INTO testcat.ns.k VALUES (1, 10.0), (2, 20.0)")
+ sql("INSERT INTO testcat.ns.m VALUES ('aa', 100.0), ('bb', 200.0)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(i, j, k, m) */ i.id, i.name, k.v, m.w
+ |FROM testcat.ns.i JOIN testcat.ns.j ON i.id = j.id AND i.name =
j.name
+ |JOIN testcat.ns.k ON i.id = k.id
+ |JOIN testcat.ns.m ON i.name = m.name
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(1, "aa", 10.0, 100.0), Row(2, "bb", 20.0,
200.0)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ with compatible transforms across three bucket
counts") {
+ val idCols = Array(Column.create("id", IntegerType))
+ Seq("b2" -> 2, "b4" -> 4, "b8" -> 8, "b12" -> 12, "c12" -> 12).foreach {
case (t, n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+ val expected = (0 until 24).map(i => Row(i, i, i))
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ // When the outer key is on the reduced side, the outer join falls
back to a shuffle.
+ Seq(
+ ("b12", "b4", "c12", 2),
+ ("b4", "b12", "c12", 0),
+ ("b8", "b4", "b2", 2),
+ ("b4", "b8", "b2", 0)).foreach { case (t1, t2, t3, numShuffles) =>
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE($t1, $t2, $t3) */ $t1.id, $t2.id, $t3.id
+ |FROM testcat.ns.$t1 JOIN testcat.ns.$t2 ON $t1.id = $t2.id
+ |JOIN testcat.ns.$t3 ON $t1.id = $t3.id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).size ==
numShuffles,
+ s"$t1 JOIN $t2 JOIN $t3 should plan $numShuffles shuffles")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ narrows common values already projected by an
inner join") {
+ val cols = Array(
+ Column.create("a", IntegerType),
+ Column.create("b", IntegerType),
+ Column.create("c", IntegerType))
+ createTable("ta", cols, Array(identity("a"), identity("b"), identity("c")))
+ createTable("tb", cols, Array(identity("b"), identity("c")))
+ createTable("tc", cols, Array.empty)
+ sql("INSERT INTO testcat.ns.ta VALUES (1, 10, 100), (2, 20, 100), (3, 30,
300)")
+ sql("INSERT INTO testcat.ns.tb VALUES (0, 10, 100), (0, 20, 100), (0, 30,
300)")
+ sql("INSERT INTO testcat.ns.tc VALUES (0, 0, 100), (0, 0, 300)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(ta, tb, tc) */ ta.a, ta.b, ta.c, tc.c
+ |FROM testcat.ns.ta JOIN testcat.ns.tb ON ta.b = tb.b AND ta.c =
tb.c
+ |JOIN testcat.ns.tc ON ta.c = tc.c
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(1, 10, 100, 100), Row(2, 20, 100, 100), Row(3,
30, 300, 300)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ keeps projected values in the order a one-side
shuffle uses") {
+ val cols = Array(
+ Column.create("store", IntegerType),
+ Column.create("dept", IntegerType),
+ Column.create("data", StringType))
+ createTable("oa", cols, Array(identity("store"), identity("dept")))
+ createTable("ob", cols, Array(identity("store"), identity("dept")))
+ createTable("oc", cols, Array.empty)
+ // The projected dept keys come out as 20, 10, which is not their sorted
order.
+ sql("INSERT INTO testcat.ns.oa VALUES (1, 20, 'a1'), (2, 10, 'a2')")
+ sql("INSERT INTO testcat.ns.ob VALUES (1, 20, 'b1'), (2, 10, 'b2')")
+ sql("INSERT INTO testcat.ns.oc VALUES (0, 10, 'c10'), (0, 20, 'c20')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(oa, ob, oc) */ oa.store, oa.data, ob.data,
oc.data
+ |FROM testcat.ns.oa JOIN testcat.ns.ob ON oa.store = ob.store AND
oa.dept = ob.dept
+ |JOIN testcat.ns.oc ON oa.dept = oc.dept
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).size == 1,
+ "only the unpartitioned oc side should shuffle")
+ checkAnswer(df, Seq(Row(1, "a1", "b1", "c20"), Row(2, "a2", "b2",
"c10")))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ does not shuffle the other side onto a reduced
scan's layout") {
+ val idCols = Array(Column.create("id", LongType))
+ Seq("l12" -> 12, "l4" -> 4, "l2" -> 2).foreach { case (t, n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+ createTable("ld", idCols, Array.empty)
+ sql("INSERT INTO testcat.ns.ld SELECT id FROM range(24)")
+ val expected = (0L until 24L).map(i => Row(i, i, i))
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ Seq("l2" -> 1, "ld" -> 2).foreach { case (third, numShuffles) =>
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(l12, l4, $third) */ l12.id, l4.id, $third.id
+ |FROM testcat.ns.l12 JOIN testcat.ns.l4 ON l12.id = l4.id
+ |JOIN testcat.ns.$third ON l12.id = $third.id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).size ==
numShuffles,
+ s"the join with $third should plan $numShuffles shuffles")
+ checkAnswer(df, expected)
+ }
+
+ // An outer join reports a single partitioning rather than a
collection.
+ val outerDf = sql(
+ """
+ |SELECT /*+ MERGE(l12, l4, ld) */ l12.id, l4.id, ld.id
+ |FROM testcat.ns.l12 LEFT JOIN testcat.ns.l4 ON l12.id = l4.id
+ |JOIN testcat.ns.ld ON l12.id = ld.id
+ |""".stripMargin)
+ checkAnswer(outerDf, expected)
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ keeps SPJ on a key the inner join did not
reduce") {
+ val cols = Array(Column.create("store", IntegerType), Column.create("id",
IntegerType))
+ createTable("ra", cols, Array(identity("store"), bucket(12, "id")))
+ createTable("rb", cols, Array(identity("store"), bucket(4, "id")))
+ createTable("rc", Array(Column.create("store", IntegerType)),
Array(identity("store")))
+ sql("INSERT INTO testcat.ns.ra SELECT CAST(id % 3 AS INT), CAST(id AS INT)
FROM range(12)")
+ sql("INSERT INTO testcat.ns.rb SELECT CAST(id % 3 AS INT), CAST(id AS INT)
FROM range(12)")
+ sql("INSERT INTO testcat.ns.rc VALUES (0), (1), (2)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(ra, rb, rc) */ ra.id, rb.id, rc.store
+ |FROM testcat.ns.ra JOIN testcat.ns.rb ON ra.store = rb.store AND
ra.id = rb.id
+ |JOIN testcat.ns.rc ON ra.store = rc.store
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "SPJ should be triggered for both joins")
+ checkAnswer(df, (0 until 12).map(i => Row(i, i, i % 3)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ does not push into a broadcast build side") {
+ createTable("qa", Array(Column.create("id", IntegerType),
Column.create("x", StringType)),
+ Array(identity("id")))
+ createTable("qd", Array(Column.create("x", StringType),
Column.create("region", IntegerType)),
+ Array(identity("region")))
+ createTable("qc", Array(Column.create("id", IntegerType)),
Array(identity("id")))
+ sql("INSERT INTO testcat.ns.qa VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+ sql("INSERT INTO testcat.ns.qd VALUES ('a', 10), ('b', 20), ('c', 30)")
+ sql("INSERT INTO testcat.ns.qc VALUES (1), (2), (3)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = sql(
+ """
+ |SELECT /*+ BROADCAST(qd) */ qa.id, qd.region, qc.id
+ |FROM testcat.ns.qa JOIN testcat.ns.qd ON qa.x = qd.x
+ |JOIN testcat.ns.qc ON qa.id = qc.id
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(1, 10, 1), Row(2, 20, 2), Row(3, 30, 3)))
+ }
+ }
+ }
+
+ test("SPARK-60232: ordering a subset-key SPJ keeps the join's key
positions") {
+ val cols = Array(Column.create("a", IntegerType), Column.create("b",
IntegerType))
+ createTable("oa3", cols, Array(identity("a"), identity("b")))
+ sql("INSERT INTO testcat.ns.oa3 VALUES (1, 10), (2, 20), (3, 10), (4, 30)")
+
+ createTable("ob3", cols, Array(identity("b")))
+ sql("INSERT INTO testcat.ns.ob3 VALUES (0, 10), (0, 20)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(oa3, ob3) */ oa3.b, oa3.a, ob3.a
+ |FROM testcat.ns.oa3 LEFT JOIN testcat.ns.ob3 ON oa3.b = ob3.b
+ |ORDER BY oa3.b
+ |""".stripMargin)
+
+ // Rows with the same b may come in either order, so check the order
on b alone.
+ val rows = df.collect().toSeq
+ assert(rows.map(_.getInt(0)) == Seq(10, 10, 20, 30), "rows should be
ordered by b")
+ assert(rows.toSet ==
+ Set(Row(10, 1, 0), Row(10, 3, 0), Row(20, 2, 0), Row(30, 4, null)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ sees a reduced scan through an alias") {
+ val idCols = Array(Column.create("id", LongType))
+ Seq("l12" -> 12, "l4" -> 4, "l3" -> 3).foreach { case (t, n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ // The nested join may also output the original column next to its
alias.
+ Seq("", "l12.id, ").foreach { extraColumn =>
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(x, l3) */ x.id2, x.id4, l3.id
+ |FROM (
+ | SELECT /*+ MERGE(l12, l4) */ ${extraColumn}l12.id AS id2,
l4.id AS id4
+ | FROM testcat.ns.l12 JOIN testcat.ns.l4 ON l12.id = l4.id) x
+ |JOIN testcat.ns.l3 ON x.id2 = l3.id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).size == 2,
+ "the outer join should fall back to shuffling both sides")
+ checkAnswer(df, (0L until 24L).map(i => Row(i, i, i)))
+ }
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ shuffles a compatible nested join built on a
reduced scan") {
+ val idCols = Array(Column.create("id", LongType))
+ Seq("nb4" -> 4, "nb12" -> 12).foreach { case (t, n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+ createTable("nc6", idCols, Array(bucket(6, "id")))
+ sql("INSERT INTO testcat.ns.nc6 VALUES (0), (1), (2), (3), (6)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(nb4, nb12, nc6) */ nb4.id, nb12.id, nc6.id
+ |FROM testcat.ns.nb4 JOIN testcat.ns.nb12 ON nb4.id = nb12.id
+ |JOIN testcat.ns.nc6 ON nb12.id = nc6.id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).size == 1,
+ "the nested join should be shuffled onto nc6's layout")
+ checkAnswer(df, Seq(0L, 1L, 2L, 3L, 6L).map(i => Row(i, i, i)))
+
+ // An outer join reports a single partitioning rather than a
collection.
+ val outerDf = sql(
+ """
+ |SELECT /*+ MERGE(nb4, nb12, nc6) */ nb12.id, nb4.id, nc6.id
+ |FROM testcat.ns.nb12 LEFT JOIN testcat.ns.nb4 ON nb12.id = nb4.id
+ |JOIN testcat.ns.nc6 ON nb12.id = nc6.id
+ |""".stripMargin)
+ assert(collectShuffles(outerDf.queryExecution.executedPlan).size == 1,
+ "the outer join below should be shuffled onto nc6's layout")
+ checkAnswer(outerDf, Seq(0L, 1L, 2L, 3L, 6L).map(i => Row(i, i, i)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ does not reshuffle a nested join above a
shuffle") {
+ val idCols = Array(Column.create("id", LongType))
+ Seq("l12" -> 12, "l4" -> 4, "c12" -> 12, "d12" -> 12).foreach { case (t,
n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ // The second join shuffles the nested one onto c12's layout. The
reduced l12 below that
+ // shuffle no longer decides the layout the third join sees.
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(l12, l4, c12, d12) */ l12.id, l4.id, c12.id,
d12.id
+ |FROM testcat.ns.l12 JOIN testcat.ns.l4 ON l12.id = l4.id
+ |JOIN testcat.ns.c12 ON l12.id = c12.id
+ |JOIN testcat.ns.d12 ON l12.id = d12.id
+ |""".stripMargin)
+ // collectShuffles would count that shuffle once for each join above
it.
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size == 1,
+ "only the nested join should be shuffled, once")
+ checkAnswer(df, (0L until 24L).map(i => Row(i, i, i, i)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ sees a reduced scan through a broadcast join's
build key") {
+ val idCols = Array(Column.create("id", LongType))
+ Seq("l12" -> 12, "l4" -> 4, "l3" -> 3).foreach { case (t, n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+ createTable("lx", idCols, Array.empty)
+ sql("INSERT INTO testcat.ns.lx SELECT id FROM range(24)")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ // The broadcast join also reports the reduced l12's partitioning on
lx.id.
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(l12, l4), BROADCAST(lx), MERGE(l3) */ l12.id,
lx.id, l3.id
+ |FROM testcat.ns.l12 JOIN testcat.ns.l4 ON l12.id = l4.id
+ |JOIN testcat.ns.lx ON l12.id = lx.id
+ |JOIN testcat.ns.l3 ON lx.id = l3.id
+ |""".stripMargin)
+ checkAnswer(df, (0L until 24L).map(i => Row(i, i, i)))
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ survives EnsureRequirements re-running for a
skewed join") {
+ val idCols = Array(Column.create("id", LongType))
+ Seq("b4" -> 4, "b12" -> 12, "b6" -> 6).foreach { case (t, n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+
+ withTempView("skewData1", "skewData2") {
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+ val skewQuery = "SELECT key1 AS k, count(*) AS c " +
+ "FROM skewData1 JOIN skewData2 ON key1 = key2 GROUP BY key1"
+ val expected = sql(skewQuery).collect().toSeq ++ (0L until 24L).map(i =>
Row(i, 1L))
+
+ // With forced skew join optimization, OptimizeSkewedJoin re-applies
EnsureRequirements to
+ // the whole plan, including the nested SPJ in the other branch that is
already planned.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key ->
"false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10",
+ SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |$skewQuery
+ |UNION ALL
+ |SELECT /*+ MERGE(b4, b12, b6) */ b4.id AS k, 1L AS c
+ |FROM testcat.ns.b4 JOIN testcat.ns.b12 ON b4.id = b12.id
+ |JOIN testcat.ns.b6 ON b4.id = b6.id
+ |""".stripMargin)
+ checkAnswer(df, expected)
+ }
+ }
+ }
+
+ test("SPARK-60232: re-running EnsureRequirements does not apply a reducer
twice") {
+ Seq(ReducibleHoursFunction, ReducibleDaysFunction,
ReducibleYearsFunction).foreach { f =>
+ catalog.createFunction(Identifier.of(Array.empty, f.name()), new
UnboundReducibleFunction(f))
+ }
+ val cols = Array(Column.create("id", LongType), Column.create("ts",
TimestampType))
+ Seq("qh" -> hours("ts"), "qd" -> days("ts"), "qy" -> years("ts")).foreach
{ case (t, p) =>
+ createTable(t, cols, Array(p))
+ sql(s"INSERT INTO testcat.ns.$t " +
+ "SELECT id, TIMESTAMP '2026-01-01 00:00:00' + id * INTERVAL '1' HOUR
FROM range(72)")
+ }
+
+ withTempView("skewData1", "skewData2") {
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+ val skewQuery = "SELECT key1 AS k, count(*) AS c " +
+ "FROM skewData1 JOIN skewData2 ON key1 = key2 GROUP BY key1"
+ val expected = sql(skewQuery).collect().toSeq ++ (0L until 72L).map(i =>
Row(i, 1L))
+
+ // The re-run's inner join must not derive the hours-to-days reducer
again and chain it after
+ // the one it pushed the first time, since floorDiv(h, 24) applied twice
is wrong.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key ->
"false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10",
+ SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |$skewQuery
+ |UNION ALL
+ |SELECT /*+ MERGE(qh, qd, qy) */ qh.id AS k, 1L AS c
+ |FROM testcat.ns.qh JOIN testcat.ns.qd ON qh.ts = qd.ts
+ |JOIN testcat.ns.qy ON qd.ts = qy.ts
+ |""".stripMargin)
+ checkAnswer(df, expected)
+ }
+ }
+ }
+
+ test("SPARK-60232: re-running EnsureRequirements keeps a single SPJ's
pairing") {
+ Seq(ReducibleHoursFunction, ReducibleDaysFunction,
ReducibleYearsFunction).foreach { f =>
+ catalog.createFunction(Identifier.of(Array.empty, f.name()), new
UnboundReducibleFunction(f))
+ }
+ val idCols = Array(Column.create("id", LongType))
+ Seq("b4" -> 4, "b12" -> 12).foreach { case (t, n) =>
+ createTable(t, idCols, Array(bucket(n, "id")))
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+ val tsCols = Array(Column.create("id", LongType), Column.create("ts",
TimestampType))
+ Seq("qh" -> hours("ts"), "qd" -> days("ts")).foreach { case (t, p) =>
+ createTable(t, tsCols, Array(p))
+ sql(s"INSERT INTO testcat.ns.$t " +
+ "SELECT id, TIMESTAMP '2026-01-01 00:00:00' + id * INTERVAL '1' HOUR
FROM range(72)")
+ }
+
+ withTempView("skewData1", "skewData2") {
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id % 3 as key1", "id as value1")
+ .createOrReplaceTempView("skewData1")
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as key2", "id as value2")
+ .createOrReplaceTempView("skewData2")
+ val skewQuery = "SELECT key1 AS k, count(*) AS c " +
+ "FROM skewData1 JOIN skewData2 ON key1 = key2 GROUP BY key1"
+ val skewRows = sql(skewQuery).collect().toSeq
+
+ // The re-run finds this join's own scans directly below it, still
carrying the reducer the
+ // first run pushed. With v2 bucketing shuffle off, as by default,
refusing that pairing
+ // would leave no spec to shuffle onto, so both sides would get a hash
shuffle.
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.ADAPTIVE_OPTIMIZE_SKEWS_IN_REBALANCE_PARTITIONS_ENABLED.key ->
"false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SKEW_JOIN_SKEWED_PARTITION_THRESHOLD.key -> "100",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "100",
+ SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10",
+ SQLConf.ADAPTIVE_FORCE_OPTIMIZE_SKEWED_JOIN.key -> "true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ Seq(
+ ("b4", "b12", "id", 24),
+ ("qh", "qd", "ts", 72)).foreach { case (t1, t2, key, numRows) =>
+ val df = sql(
+ s"""
+ |$skewQuery
+ |UNION ALL
+ |SELECT /*+ MERGE($t1, $t2) */ $t1.id AS k, 1L AS c
+ |FROM testcat.ns.$t1 JOIN testcat.ns.$t2 ON $t1.$key = $t2.$key
+ |""".stripMargin)
+ checkAnswer(df, skewRows ++ (0L until numRows.toLong).map(i =>
Row(i, 1L)))
+ // Only the skewed branch shuffles: both sides of its join, and its
aggregate.
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size == 3,
+ s"$t1 JOIN $t2 should keep SPJ when EnsureRequirements re-runs")
+ }
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ falls back when a broadcast build key cannot
be traced") {
+ val longCols = Array(Column.create("id", LongType))
+ val intCols = Array(Column.create("id", IntegerType))
+ Seq("l12" -> 12, "l4" -> 4, "l3" -> 3).foreach { case (t, n) =>
+ createTable(t, longCols, Array(bucket(n, "id")))
+ }
+ Seq("lx", "ly").foreach(t => createTable(t, longCols, Array.empty))
+ createTable("lxi", intCols, Array.empty)
+ createTable("l3i", intCols, Array(bucket(3, "id")))
+ Seq("l12", "l4", "l3", "lx", "ly", "lxi", "l3i").foreach { t =>
+ sql(s"INSERT INTO testcat.ns.$t SELECT id FROM range(24)")
+ }
+
+ // Each query reaches the reduced l12 through a broadcast build key that
maps back to no
+ // single unreduced position, so the outer join has to shuffle.
+ val longRows = (0L until 24L).map(i => Row(i, i, i))
+ val queries = Seq(
+ // A build key equated to two streamed keys.
+ """
+ |SELECT /*+ MERGE(l12, l4), BROADCAST(lx), MERGE(l3) */ l12.id, lx.id,
l3.id
+ |FROM testcat.ns.l12 JOIN testcat.ns.l4 ON l12.id = l4.id
+ |JOIN testcat.ns.lx ON lx.id = l12.id AND lx.id = l4.id
+ |JOIN testcat.ns.l3 ON lx.id = l3.id
+ |""".stripMargin -> longRows,
+ // A cast build key.
+ """
+ |SELECT /*+ MERGE(l12, l4), BROADCAST(lxi), MERGE(l3i) */ l12.id,
lxi.id, l3i.id
+ |FROM testcat.ns.l12 JOIN testcat.ns.l4 ON l12.id = l4.id
+ |JOIN testcat.ns.lxi ON l12.id = lxi.id
+ |JOIN testcat.ns.l3i ON lxi.id = l3i.id
+ |""".stripMargin -> (0 until 24).map(i => Row(i.toLong, i, i)),
+ // A build key that a broadcast inside the build side also uses.
+ """
+ |SELECT /*+ MERGE(l12, l4), BROADCAST(b), MERGE(l3) */ l12.id, b.id,
l3.id
+ |FROM testcat.ns.l12 JOIN testcat.ns.l4 ON l12.id = l4.id
+ |JOIN (
+ | SELECT /*+ BROADCAST(ly) */ lx.id AS xid, ly.id
+ | FROM testcat.ns.lx JOIN testcat.ns.ly ON lx.id = ly.id) b ON l12.id
= b.id
+ |JOIN testcat.ns.l3 ON b.id = l3.id
+ |""".stripMargin -> longRows)
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ queries.foreach { case (query, expected) =>
+ val df = sql(query)
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size == 2,
+ "the outer join should fall back to shuffling both sides")
+ checkAnswer(df, expected)
+ }
+ }
+ }
+ }
+
+ test("SPARK-60232: nested SPJ keeps SPJ on an unreduced column with the same
transform") {
+ val cols = Array(Column.create("a", LongType), Column.create("b",
LongType))
+ createTable("pa", cols, Array(bucket(12, "a"), bucket(12, "b")))
+ createTable("pb", cols, Array(bucket(4, "a"), bucket(12, "b")))
+ createTable("pc", cols, Array(bucket(12, "b")))
+ Seq("pa", "pb", "pc").foreach { t =>
+ sql(s"INSERT INTO testcat.ns.$t SELECT id, id FROM range(24)")
+ }
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ // The inner join reduces pa at a only, so bucket(12, b) still holds
pa's own values.
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(pa, pb, pc) */ pa.a, pc.b
+ |FROM testcat.ns.pa JOIN testcat.ns.pb ON pa.a = pb.a AND pa.b =
pb.b
+ |JOIN testcat.ns.pc ON pa.b = pc.b
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "SPJ should be triggered for both joins")
+ checkAnswer(df, (0L until 24L).map(i => Row(i, i)))
+ }
+ }
+ }
+
+ test("SPARK-60232: ordering a subset-key SPJ range-shuffles when a one-side
shuffle is below") {
+ val cols = Array(
+ Column.create("store", IntegerType),
+ Column.create("dept", IntegerType),
+ Column.create("data", StringType))
+ createTable("sa", cols, Array(identity("store"), identity("dept")))
+ // The dept keys come out as 20, 10, so ordering by dept reorders them.
+ sql("INSERT INTO testcat.ns.sa VALUES (1, 20, 'a1'), (2, 10, 'a2')")
+
+ createTable("sb", cols, Array.empty)
+ sql("INSERT INTO testcat.ns.sb VALUES (0, 10, 'b10'), (0, 20, 'b20')")
+
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString,
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true",
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") {
+ Seq("LEFT JOIN", "JOIN").foreach { joinType =>
+ val df = sql(
+ s"""
+ |SELECT /*+ MERGE(sa, sb) */ sa.dept, sa.data, sb.data, sa.store
+ |FROM testcat.ns.sa $joinType testcat.ns.sb ON sa.dept = sb.dept
+ |ORDER BY sa.dept
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(10, "a2", "b10", 2), Row(20, "a1", "b20",
1)))
+ }
+ }
+ }
+ }
+
+ test("SPARK-60232: SPJ does not push common values below a shuffle") {
Review Comment:
**Finding 43.** This test covers SPARK-60552, and its `orderedDf` is that
ticket's LEFT JOIN query. Naming it after SPARK-60552 lets a search by that
ticket find it:
```suggestion
test("SPARK-60552: SPJ does not push common values below a shuffle") {
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -502,6 +537,15 @@ 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 38.** The order I suggested in [finding
34](https://github.com/apache/spark/pull/59599#discussion_r4412684753) refuses
a reduced spec before it looks for an earlier run's pairing. So on the re-run,
the nested skew query's outer join still refuses the reducer it pushed into the
inner scans the first time. Its 5 shuffles are 2 more than `branch-4.1` plans.
Checking for the earlier pairing first keeps it there too:
```scala
// 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))
}
if (Seq(left, right).zip(specs).exists { case (p, s) =>
isBuiltFromReducedScan(p, s) }) {
return None
}
```
This only changes re-runs. On a first run no child carries a reducer outside
a join, because only this method pushes reducers, and the join it plans then
sits above them. On a re-run such a reducer comes from this join's own SPJ, or
from an outer one that pushed through both of its sides, so the earlier pairing
still holds.
What I ran: `<skewed query> UNION ALL <nested SPJ>`, with the skewed query
and the confs of the skew tests, AQE on.
| nested SPJ next to the skewed join | `branch-4.1` | this head | with the
swap |
|---|---|---|---|
| the query of `nested SPJ survives EnsureRequirements re-running for a
skewed join` | 3 | 5 | 3 |
| the same with `ON b12.id = b6.id` | 3, with 8 of 24 SPJ rows | 4 | 4 |
| `qh JOIN qd JOIN qy` of `re-running EnsureRequirements does not apply a
reducer twice` | 0 of 72 SPJ rows | 4 | 4 |
The other cells return the right rows. In the last two rows the outer join's
refusal stands, because its third side carries no reducer. With the swap, the
first test could assert its 3 shuffles. The description's count of 5 would
change to 3.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -199,8 +212,14 @@ case class EnsureRequirements(
// from one member, otherwise the sides end up grouped on different
keys, or on a key set the
// child does not even have. Pick that member once, preferring the
finest when several
// qualify. Only the branch that shuffles a child reads these, hence
`lazy`.
+ // A child with a member built from a reduced scan only looks compatible
on the transform
+ // the scan still reports, so it is shuffled onto the best spec instead.
lazy val matchedIndexes = bestSpecOpt.toSeq.flatMap { best =>
- childrenIndexes.filter(i => best.isCompatibleWith(specs(i)))
+ childrenIndexes.filter(i => best.isCompatibleWith(specs(i)) &&
+ !flattenSpec(specs(i)).exists {
+ case k: KeyGroupedShuffleSpec =>
isBuiltFromReducedScan(children(i), k)
Review Comment:
**Finding 42.** This is the negation of the member check in `candidateSpecs`
at lines 157-160. One helper next to `isBuiltFromReducedScan`, e.g.
`hasMemberBuiltFromReducedScan(child, spec)`, would keep the two from drifting
apart.
--
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]
