peter-toth commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4347644825
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -144,7 +145,9 @@ case class EnsureRequirements(
// To choose a KeyGroupedShuffleSpec, we must be able to push down
SPJ parameters into
// the scan (for join key positions). If these parameters can't be
pushed down, this
// spec can't be used to shuffle other children.
- case (idx, _: KeyGroupedShuffleSpec) =>
canPushDownSPJParamsToScan(children(idx))
+ case (idx, s: KeyGroupedShuffleSpec) =>
+ canPushDownSPJParamsToScan(children(idx)) &&
+ !isBuiltFromReducedScan(children(idx), s)
Review Comment:
**Finding 14.** This check covers a child that reports a single
partitioning. A nested inner join reports a `PartitioningCollection`, so its
spec is a `ShuffleSpecCollection`, and it takes the `case _ => true` branch
below. When `checkKeyGroupCompatible` refuses a spec built from the reduced
scan, the planner falls back to this path, and the reduced scan's member can
still become `bestMemberOpt`. The other side is then shuffled onto `bucket(12,
id)` with values in the `bucket(4, id)` domain.
With BIGINT ids 0 to 23 in `bucket(12)`, `bucket(4)` and `bucket(2)` tables,
`allowCompatibleTransforms` and `shuffle.enabled`, AQE on and off:
- `l12 JOIN l4 JOIN l2`, both joins on `l12.id`: 16 of 24 rows here.
`branch-4.1` returns all 24 with no shuffle.
- `l12 JOIN l4 JOIN ld`, with `ld` unpartitioned: 16 of 24 rows, here and on
`branch-4.1`.
Checking the members of a collection too fixes both:
```scala
case (idx, s: ShuffleSpecCollection) =>
flattenSpec(s).forall {
case k: KeyGroupedShuffleSpec =>
!isBuiltFromReducedScan(children(idx), k)
case _ => true
}
```
With it, the two queries return all 24 rows, with one and two shuffles, and
`KeyGroupedPartitioningSuite` and `EnsureRequirementsSuite` pass. None of the
new tests combines `allowCompatibleTransforms` with `shuffle.enabled`, so the
first query would make a good test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,7 +709,121 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition information down to the scan nodes
+ /**
+ * Compose join-key positions from an outer SPJ with positions already
pushed by an inner SPJ.
+ * Outer positions index the layout the outer join sees; inner positions
index the scan's raw
+ * partition expressions, so they are composed rather than replaced.
+ */
+ private def composeJoinKeyPositions(
+ existing: Option[Seq[Int]],
+ newPositions: Option[Seq[Int]]): Option[Seq[Int]] = {
+ (existing, newPositions) match {
+ case (Some(old), Some(newPos)) => Some(newPos.map(old))
+ case (_, None) => existing
+ case (None, Some(newPos)) => Some(newPos)
+ }
+ }
+
+ /**
+ * Project common partition values onto `projectionPositions` of the layout
`expressions`
+ * describes and deduplicate. Values that collapse to the same key are read
as one group, so
+ * they become one partition, unless the scan is partially clustered, where
each split of the
+ * group still gets its own partition and the split counts add up.
+ */
+ private def projectAndDedupeCommonPartitionValues(
+ values: Seq[(InternalRow, Int)],
+ expressions: Seq[Expression],
+ projectionPositions: Seq[Int],
+ partiallyClustered: Boolean): Seq[(InternalRow, Int)] = {
+ val projectedExpressions = projectionPositions.map(expressions(_))
+ val partitionOrdering = RowOrdering.createNaturalAscendingOrdering(
+ projectedExpressions.map(_.dataType))
+ values.map { case (partValue, numSplits) =>
+ val projected = KeyGroupedPartitioning.project(
+ expressions, projectionPositions, partValue)
+ (InternalRowComparableWrapper(projected, projectedExpressions),
numSplits)
+ }.groupBy(_._1).map { case (wrapper, entries) =>
+ (wrapper.row, if (partiallyClustered) entries.map(_._2).sum else 1)
+ }.toSeq.sortBy(_._1)(partitionOrdering)
Review Comment:
**Finding 13.** `KeyGroupedPartitioning.apply` projects partition values
with `.distinct`, so it keeps the order in which the keys first appear. That is
the order a one-side shuffle routes the other child into, through
`createShuffleSpec` and `createPartitioning`. This method sorts the projected
values instead. When the projected key is not the leading partition key, the
two orders can differ, and partition i of each side holds a different key.
`oa` and `ob` are `(store INT, dept INT, data STRING) PARTITIONED BY (store,
dept)`, both holding `(1, 20)` and `(2, 10)`. `oc` is unpartitioned and holds
depts 10 and 20. With `shuffle.enabled`,
`allowJoinKeysSubsetOfPartitionKeys.enabled=true` and
`requireAllClusterKeysForCoPartition=false`:
```sql
SELECT /*+ MERGE(oa, ob, oc) */ oa.store, oa.data, ob.data, oc.data
FROM oa JOIN ob ON oa.store = ob.store AND oa.dept = ob.dept
JOIN oc ON oa.dept = oc.dept
```
It returns no rows instead of 2, AQE on and off. The two scans report `[10],
[20]`, while `oc` is shuffled into `keygroupedpartitioning(dept, 2, [20],
[10])`. `branch-4.1` fails the same query at planning with
`ArrayIndexOutOfBoundsException`. (`oa.store` is selected so that the nested
join keeps reporting its partitioning.) The "narrows common values" test passes
only because its projected values come out sorted already.
Keeping the first-occurrence order fixes it. With this, the query returns
both rows with one shuffle, and `KeyGroupedPartitioningSuite` and
`EnsureRequirementsSuite` pass:
```scala
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)
}
```
Could you add the query above as a test, asserting the rows and the single
shuffle?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -721,19 +846,42 @@ case class EnsureRequirements(
child, values, joinKeyPositions, reducers, applyPartialClustering,
replicatePartitions))
}
-
- private def populateJoinKeyPositions(
+ /**
+ * 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
+ }
)
)
case node =>
- node.mapChildren(child => populateJoinKeyPositions(
- child, joinKeyPositions))
+ // Recurse through every child down to the scans.
Review Comment:
**Finding 17.** This also descends through a `BroadcastExchangeExec`. A
broadcast build side does not feed this join's partitioning, but its scan still
gets the outer join's values.
`qa (id INT, x STRING) PARTITIONED BY (id)` holds ids 1 to 3. `qd (x STRING,
region INT) PARTITIONED BY (region)` holds regions 10, 20 and 30. `qc (id INT)
PARTITIONED BY (id)` holds ids 1 to 3:
```sql
SELECT /*+ BROADCAST(qd) */ qa.id, qd.region, qc.id
FROM qa JOIN qd ON qa.x = qd.x
JOIN qc ON qa.id = qc.id
```
It returns no rows instead of 3, AQE on and off, here and on `branch-4.1`
alike. `qd`'s scan gets the ids as its common partition values, so every region
partition is dropped. So this is not a regression, but it is the same push, and
stopping it at the broadcast fixes it:
```scala
case b: BroadcastExchangeLike => b
```
With it, the query returns all 3 rows, and `KeyGroupedPartitioningSuite` and
`EnsureRequirementsSuite` pass. A scan below an inner one-side shuffle fails
too, with `Can't zip RDDs` or a `PartitioningCollection` error. Stopping at
that exchange does not fix it, since the exchange keeps the partition count it
was planned with. I filed SPARK-60460 for that one.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -2980,6 +2980,512 @@ 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)
+ }
+ }
+ }
+
+ // Aligning partition counts when one id maps to multiple names needs
GroupPartitionsExec
+ // (4.2+, SPARK-55535). Projecting common values in populateSpjParams covers
the other shapes.
+ ignore("SPARK-60232: nested SPJ with multiple names per id over common
partition values") {
Review Comment:
**Finding 16.** This test passes on this head, AQE on and off, and fails on
`branch-4.1`. It is the ticket's shape 3 with two names per id, so I'd enable
it. The comment above it no longer holds, since `populateSpjParams` now
projects and de-duplicates the common values without `GroupPartitionsExec`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,7 +709,121 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition information down to the scan nodes
+ /**
+ * Compose join-key positions from an outer SPJ with positions already
pushed by an inner SPJ.
+ * Outer positions index the layout the outer join sees; inner positions
index the scan's raw
+ * partition expressions, so they are composed rather than replaced.
+ */
+ private def composeJoinKeyPositions(
+ existing: Option[Seq[Int]],
+ newPositions: Option[Seq[Int]]): Option[Seq[Int]] = {
+ (existing, newPositions) match {
+ case (Some(old), Some(newPos)) => Some(newPos.map(old))
+ case (_, None) => existing
+ case (None, Some(newPos)) => Some(newPos)
+ }
+ }
+
+ /**
+ * Project common partition values onto `projectionPositions` of the layout
`expressions`
+ * describes and deduplicate. Values that collapse to the same key are read
as one group, so
+ * they become one partition, unless the scan is partially clustered, where
each split of the
+ * group still gets its own partition and the split counts add up.
+ */
+ private def projectAndDedupeCommonPartitionValues(
+ values: Seq[(InternalRow, Int)],
+ expressions: Seq[Expression],
+ projectionPositions: Seq[Int],
+ partiallyClustered: Boolean): Seq[(InternalRow, Int)] = {
+ val projectedExpressions = projectionPositions.map(expressions(_))
+ val partitionOrdering = RowOrdering.createNaturalAscendingOrdering(
+ projectedExpressions.map(_.dataType))
+ values.map { case (partValue, numSplits) =>
+ val projected = KeyGroupedPartitioning.project(
+ expressions, projectionPositions, partValue)
+ (InternalRowComparableWrapper(projected, projectedExpressions),
numSplits)
+ }.groupBy(_._1).map { case (wrapper, entries) =>
+ (wrapper.row, if (partiallyClustered) entries.map(_._2).sum else 1)
+ }.toSeq.sortBy(_._1)(partitionOrdering)
+ }
+
+ /**
+ * The partition expressions a scan currently reports, i.e. its own
partition expressions
+ * projected onto the join-key positions pushed so far.
+ */
+ private def reportedLayout(params: StoragePartitionJoinParams):
Option[Seq[Expression]] =
+ params.keyGroupedPartitioning.map { expressions =>
+ params.joinKeyPositions.map(_.map(expressions)).getOrElse(expressions)
+ }
+
+ /**
+ * Common partition values for a scan after a push. Values pushed by this
join are the merge with
+ * the other side and already sit in the layout `newPositions` select, so
they win. Otherwise the
+ * values an inner join pushed sit in the scan's reported layout, and are
projected onto
+ * `newPositions`, which index that layout.
+ */
+ private def finalizeCommonPartitionValues(
+ existing: StoragePartitionJoinParams,
+ pushedValues: Option[Seq[(InternalRow, Int)]],
+ newPositions: Option[Seq[Int]]): Option[Seq[(InternalRow, Int)]] = {
+ (pushedValues, existing.commonPartitionValues, newPositions,
reportedLayout(existing)) match {
+ case (Some(values), _, _, _) => Some(values)
+ case (None, Some(values), Some(positions), Some(layout))
+ if values.nonEmpty && values.head._1.numFields == layout.length &&
+ positions != layout.indices =>
+ Some(projectAndDedupeCommonPartitionValues(
+ values, layout, positions, existing.applyPartialClustering))
+ case (None, values, _, _) => values
+ }
+ }
+
+ /**
+ * Reducers for a scan after a push. The reducers an inner join pushed index
the scan's reported
+ * layout, so they are first projected onto `newPositions`. A reducer this
join pushes for a
+ * position applies to values the inner reducer already produced, so the two
are chained.
+ */
+ private def composeReducers(
+ existing: Option[Seq[Option[Reducer[_, _]]]],
+ newPositions: Option[Seq[Int]],
+ pushed: Option[Seq[Option[Reducer[_, _]]]]):
Option[Seq[Option[Reducer[_, _]]]] = {
+ val projectedExisting = (existing, newPositions) match {
+ case (Some(rs), Some(positions)) => Some(positions.map(rs))
+ case _ => existing
+ }
+ (projectedExisting, pushed) match {
+ case (Some(inner), Some(outer)) =>
+ assert(inner.length == outer.length,
+ "pushed reducers must index the same layout as the composed join-key
positions")
+ Some(inner.zip(outer).map {
+ case (Some(i), Some(o)) => Some(ChainedReducer(i, o))
+ case (i, o) => o.orElse(i)
+ })
+ case (inner, outer) => outer.orElse(inner)
+ }
+ }
+
+ /**
+ * Whether `spec` was built from the reported partitioning of a scan below
`plan` whose splits an
+ * inner join already reduced. Such a scan reports its own partition
transform but values of the
+ * reduced transform, so a spec built from it describes a layout the scan
does not have.
+ */
+ 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 if
scan.spjParams.reducers.exists(_.exists(_.isDefined)) =>
+ reportedLayout(scan.spjParams).exists { layout =>
+ specExpressions.forall(e => layout.exists(_.semanticEquals(e)))
+ }
Review Comment:
**Finding 15.** This refuses the spec when the scan has a reducer at any
position, even when the spec uses only positions that are not reduced. The
values at those positions are still the scan's own, so the spec is sound there.
`ra` is `PARTITIONED BY (store, bucket(12, id))` and `rb` is `PARTITIONED BY
(store, bucket(4, id))`, both holding `(id % 3, id)` for ids 0 to 11. `rc` is
`PARTITIONED BY (store)`. With `allowCompatibleTransforms`,
`allowJoinKeysSubsetOfPartitionKeys.enabled=true` and
`requireAllClusterKeysForCoPartition=false`:
```sql
SELECT /*+ MERGE(ra, rb, rc) */ ra.id, rb.id, rc.store
FROM ra JOIN rb ON ra.store = rb.store AND ra.id = rb.id
JOIN rc ON ra.store = rc.store
```
`branch-4.1` plans it with no shuffle and returns the right rows. Here both
sides of the outer join are shuffled, AQE on and off. Checking only the
positions the spec uses brings back the plan with no shuffle, and the
three-bucket test still passes:
```suggestion
case scan: BatchScanExec =>
(reportedLayout(scan.spjParams), scan.spjParams.reducers) match {
case (Some(layout), Some(reducers)) =>
layout.zip(reducers).exists { case (e, r) =>
r.isDefined && specExpressions.exists(_.semanticEquals(e))
}
case _ => false
}
```
A no-shuffle assertion on the query above would pin it. It fails on this
head and passes with the change.
When the outer key is on the reduced side itself, as in `b12 JOIN b4 JOIN
b2` on `b12.id`, refusing is the safe choice. `branch-4.1` gets those rows
right only because a bucket reducer gives the same answer on already reduced
values. That case now costs two shuffles, which is worth a line in the
description.
--
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]
