ammarchalifah commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4377513935
##########
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:
Addressed
##########
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:
Addressed
##########
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:
Addressed
--
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]
