peter-toth commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4377309444
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +723,174 @@ 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 = {
+ 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) =>
+ // A `Project` between the joins can rename the column under
the transform, so
+ // match a transform by its function rather than by the column
it reads.
+ r.isDefined && specExpressions.exists {
+ case s: TransformExpression => e match {
+ case t: TransformExpression => s.isSameFunction(t)
Review Comment:
**Finding 22.** [Finding
19](https://github.com/apache/spark/pull/59599#discussion_r4359560138) noted
that matching by function alone can shuffle a side that is fine. Here is that
shape.
`pa` is `PARTITIONED BY (bucket(12, a), bucket(12, b))`, `pb` is
`PARTITIONED BY (bucket(4, a), bucket(12, b))` and `pc` is `PARTITIONED BY
(bucket(12, b))`. Each holds `(id, id)` for BIGINT ids 0 to 23. With
`allowCompatibleTransforms`, `allowJoinKeysSubsetOfPartitionKeys.enabled=true`
and `requireAllClusterKeysForCoPartition=false`:
```sql
SELECT /*+ MERGE(pa, pb, pc) */ pa.a, pc.b
FROM pa JOIN pb ON pa.a = pb.a AND pa.b = pb.b
JOIN pc ON pa.b = pc.b
```
The inner join reduces `pa` at `a` only. The outer join uses `b`, but
`bucket(12, b)` is the same function as the reduced `bucket(12, a)`, so the
spec is refused. `branch-4.1` returns the 24 rows with no shuffle, and this
head returns them with two, AQE on and off.
Comparing the columns again has to see through the alias. Checking whether
the scan's column is still in the child's output is not enough. When the nested
join outputs both `l12.id` and `l12.id AS id2`, the alias test's query on `id2`
returns 8 of 24 rows that way. Mapping the spec's columns back through the
aliases first, and then comparing them exactly, handles both:
```scala
// 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.
val aliases = AttributeMap(plan.flatMap(_.expressions).flatMap(_.collect
{
case a @ Alias(child: Attribute, _) => a.toAttribute -> child
}))
def unalias(e: Expression): Expression = e.transformUp {
case a: Attribute if aliases.contains(a) => unalias(aliases(a))
}
val specExpressions = spec.partitioning.expressions.map(unalias)
...
r.isDefined && specExpressions.exists(_.semanticEquals(e))
```
With it, the query above plans no shuffle. The alias test's query returns 24
rows, with and without `l12.id` next to the alias. All 20 SPARK-60232 tests
pass with it, and both suites pass with it together with the fixes in 23 and 24.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -345,10 +362,11 @@ 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)
+ populateSpjParams(plan, commonPartitionValues = Some(sorted.map((_,
1))),
Review Comment:
**Finding 23.** A one-side shuffle routes its other child into the order the
partition values had when the join was planned. With a subset of the partition
keys, that is the first-appearance order of `KeyGroupedPartitioning.apply`.
This push then sorts the scans' values, but the shuffle keeps the old order, so
partition i of each side can hold a different key.
`sa` is `(store INT, dept INT, data STRING) PARTITIONED BY (store, dept)`
and holds `(1, 20, 'a1')` and `(2, 10, 'a2')`. `sb` has the same columns, is
not partitioned, and holds `(0, 10, 'b10')` and `(0, 20, 'b20')`. With
`shuffle.enabled`, `bucketing.sorting.enabled`,
`allowJoinKeysSubsetOfPartitionKeys.enabled=true` and
`requireAllClusterKeysForCoPartition=false`:
```sql
SELECT /*+ MERGE(sa, sb) */ sa.dept, sa.data, sb.data, sa.store
FROM sa LEFT JOIN sb ON sa.dept = sb.dept
ORDER BY sa.dept
```
- Here the LEFT JOIN returns `[10,a2,null,2]` and `[20,a1,null,1]`, and the
same query with an inner `JOIN` returns no rows, AQE on and off. `sb` is
shuffled into `keygroupedpartitioning(dept, 2, [20], [10])`, while the scan is
sorted to `[10], [20]`.
- On `branch-4.1`, both fail at planning with
`ArrayIndexOutOfBoundsException: Index 1 out of bounds for length 1` with AQE
on, and return no rows with AQE off.
- Without the `ORDER BY`, or with sorting off, the LEFT JOIN returns the
right rows here.
So with AQE on, keeping the positions here, which finding 18 asked for,
turns a planning failure into wrong rows. With AQE off, it leaves the wrong
rows in place. (`sa.store` is selected so that the scan keeps reporting its
partitioning.)
Falling back to a range shuffle fixes it when such a shuffle sits below and
the sort changes the order. With AQE the shuffle can already be a
`ShuffleQueryStageExec`, so the check looks for that too:
```scala
val sorted = partitionValues.sorted(partitionOrdering)
// 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 {
case s @ (_: ShuffleExchangeLike | _: ShuffleQueryStageExec) =>
s.outputPartitioning.isInstanceOf[KeyGroupedPartitioning]
case _ => false
}
val reorders = partitionValues.zip(sorted).exists { case (a, b) =>
partitionOrdering.compare(a, b) != 0
}
if (hasKeyGroupedShuffle && reorders) {
ShuffleExchangeExec(d.createPartitioning(conf.numShufflePartitions), plan)
} else {
populateSpjParams(plan, commonPartitionValues =
Some(sorted.map((_, 1))),
joinKeyPositions = None)
}
```
It needs `import
org.apache.spark.sql.execution.adaptive.ShuffleQueryStageExec`. With it, both
queries return the right rows with a range shuffle, AQE on and off.
`KeyGroupedPartitioningSuite` and `EnsureRequirementsSuite` pass. `reorders`
keeps the plan unchanged when the values are already in order, as with a single
identity key. The query above would make a good test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -199,8 +209,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 whose matching member comes from a reduced scan only looks
compatible on the
+ // transform the scan still reports, so it is shuffled onto the best
spec instead.
Review Comment:
**Finding 26.** The check below shuffles the child when any member of its
spec comes from a reduced scan, not only the matching one.
```suggestion
// 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.
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +723,174 @@ 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 = {
+ 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) =>
+ // A `Project` between the joins can rename the column under
the transform, so
+ // match a transform by its function rather than by the column
it reads.
+ r.isDefined && specExpressions.exists {
+ case s: TransformExpression => e match {
+ case t: TransformExpression => s.isSameFunction(t)
+ case _ => false
+ }
+ case s => s.semanticEquals(e)
+ }
+ }
+ case _ => false
+ }
+ case _ => false
+ }
+ case _ => false
+ }
+ }
+
+ /**
+ * Push SPJ parameters down to every scan below `plan`, composing join-key
positions with any
+ * already pushed by an inner join instead of replacing them.
+ */
+ private def populateSpjParams(
plan: SparkPlan,
- joinKeyPositions: Option[Seq[Int]]): SparkPlan = plan match {
+ commonPartitionValues: Option[Seq[(InternalRow, Int)]],
+ joinKeyPositions: Option[Seq[Int]],
+ reducers: Option[Seq[Option[Reducer[_, _]]]] = None,
+ applyPartialClustering: Boolean = false,
+ replicatePartitions: Boolean = false): SparkPlan = plan match {
case scan: BatchScanExec =>
+ val existing = scan.spjParams
+ val composedPositions = composeJoinKeyPositions(
+ existing.joinKeyPositions, joinKeyPositions)
scan.copy(
- spjParams = scan.spjParams.copy(
- joinKeyPositions = joinKeyPositions
+ spjParams = existing.copy(
+ commonPartitionValues = finalizeCommonPartitionValues(
+ existing, commonPartitionValues, joinKeyPositions),
+ joinKeyPositions = composedPositions,
+ reducers = composeReducers(existing.reducers, joinKeyPositions,
reducers),
+ applyPartialClustering = commonPartitionValues match {
+ case Some(_) => applyPartialClustering
+ case _ => existing.applyPartialClustering
+ },
+ replicatePartitions = commonPartitionValues match {
+ case Some(_) => replicatePartitions
+ case _ => existing.replicatePartitions
+ }
)
)
+ // A broadcast build side does not feed this join's partitioning.
+ case b: BroadcastExchangeLike => b
Review Comment:
**Finding 24.** A shuffle's input does not feed this join's partitioning
either, but the push still reaches the scans below it. A partitioned scan there
gets the other side's values and drops every partition whose key is not among
them. Two shapes, both with `shuffle.enabled`, return the wrong rows here and
on `branch-4.1`, AQE on and off:
- **Non-nested.** `sa7 (dept INT) PARTITIONED BY (dept)` holds depts 1 to 3.
`sb7 (dept INT, region INT) PARTITIONED BY (region)` holds `(1, 10)`, `(2, 20)`
and `(3, 30)`. With `bucketing.sorting.enabled`, `SELECT /*+ MERGE(sa7, sb7) */
sa7.dept, sb7.region FROM sa7 LEFT JOIN sb7 ON sa7.dept = sb7.dept ORDER BY
sa7.dept` returns `region` as null on every row. `sb7` is shuffled onto `sa7`'s
layout, and `ensureOrdering` then pushes `sa7`'s values into `sb7`'s scan.
- **Nested.** `pa (id INT)` and `pc (id INT)` are `PARTITIONED BY (id)` and
hold ids 1 to 3. `pb (id INT, region INT) PARTITIONED BY (region)` holds `(1,
10)`, `(2, 20)` and `(3, 30)`. `SELECT /*+ MERGE(pa, pb, pc) */ pa.id,
pb.region, pc.id FROM pa JOIN pb ON pa.id = pb.id JOIN pc ON pa.id = pc.id`
returns no rows instead of 3. The outer join's push reaches `pb`'s scan below
the inner shuffle.
Stopping the push at the shuffle fixes both:
```suggestion
case b: BroadcastExchangeLike => b
// Nor does the input of a shuffle.
case s: ShuffleExchangeLike => s
```
With it, both queries return the right rows, and the two suites pass.
[Finding 17](https://github.com/apache/spark/pull/59599#discussion_r4347644857)
said stopping at the exchange does not fix the nested one-side shuffle. That
still holds for the partition count, which
[SPARK-60460](https://issues.apache.org/jira/browse/SPARK-60460) tracks. It
does fix the dropped rows, though. I filed
[SPARK-60552](https://issues.apache.org/jira/browse/SPARK-60552) for the
non-nested shape.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -2980,6 +2980,708 @@ 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") {
+ Seq(
+ ("b12", "b4", "c12"),
+ ("b4", "b12", "c12"),
+ ("b8", "b4", "b2"),
+ ("b4", "b8", "b2")).foreach { case (t1, t2, t3) =>
+ 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)
+ checkAnswer(df, expected)
Review Comment:
**Finding 25.** This test and the fallback ones (`does not shuffle the other
side onto a reduced scan's layout`, `sees a reduced scan through an alias`,
`shuffles a compatible nested join built on a reduced scan`) check the rows
only. Three additions would pin the fallback:
- **Shuffle counts here.** Two shuffles with `b12` or `b8` first, none with
`b4` first, AQE on and off.
- **22's query.** The same transform on two columns pins that a position the
inner join did not reduce stays eligible.
- **An outer join below.** With the reduced-layout test's tables, `l12 LEFT
JOIN l4 ON l12.id = l4.id JOIN ld ON l12.id = ld.id` returns 16 of 24 rows on
`branch-4.1` and all 24 here. An outer join reports a single partitioning, so
its spec takes the `KeyGroupedShuffleSpec` case of `candidateSpecs` rather than
the collection one, and no test goes 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]
