ammarchalifah commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4377520152
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,42 +709,131 @@ case class EnsureRequirements(
joinType == LeftAnti || joinType == LeftOuter
}
- // Populate the common partition information down to the scan nodes
- private def populateCommonPartitionInfo(
- plan: SparkPlan,
- values: Seq[(InternalRow, Int)],
+ /**
+ * 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)
+ }
+ }
+
+ /** Controls which SPJ fields are updated on a scan during pushdown. */
+ private case class SpjPushFlags(
+ pushJoinKeyPositions: Boolean = true,
+ pushCommonValues: Boolean = true)
+
+ /**
+ * Compute side-aware SPJ push flags for nested shuffle joins. With no
join-key positions,
+ * push common values to both inner sides. With subset join keys, compose
join-key positions
+ * on the matching side only and retain inner common-partition metadata on
both sides.
+ */
+ private def nestedSpjPushFlags(
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))
+ isMatchingSide: Boolean): SpjPushFlags = {
+ val syncCommonValuesToBothSides = joinKeyPositions.isEmpty
+ if (joinKeyPositions.isDefined && !syncCommonValuesToBothSides) {
+ SpjPushFlags(pushJoinKeyPositions = isMatchingSide, pushCommonValues =
false)
+ } else {
+ SpjPushFlags(
+ pushJoinKeyPositions = isMatchingSide || syncCommonValuesToBothSides,
+ pushCommonValues = isMatchingSide || syncCommonValuesToBothSides)
+ }
}
+ /**
+ * Push SPJ parameters down to the scan in each child. When a child is
itself a join, push to
+ * both inner sides so partition counts stay aligned, but only compose
join-key positions on
+ * the side that feeds this join. When join keys are a subset of partition
keys, keep the inner
+ * join's common partition values on both sides and only push join-key
positions to the
+ * matching side.
+ */
+ private def populateSpjParamsOnSide(
+ plan: SparkPlan,
+ fromLeft: Boolean,
+ 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 smj: SortMergeJoinExec =>
+ // Decide per side whether to push join-key positions, common values, or
retain inner
+ // SPJ metadata already pushed by a nested join below.
+ val leftFlags = nestedSpjPushFlags(joinKeyPositions, isMatchingSide =
fromLeft)
+ val rightFlags = nestedSpjPushFlags(joinKeyPositions, isMatchingSide =
!fromLeft)
+ smj.copy(
+ // Pass join-key positions only to the side that feeds this join.
+ left = populateSpjParams(
+ smj.left, commonPartitionValues, if (fromLeft) joinKeyPositions else
None,
+ reducers, applyPartialClustering, replicatePartitions, leftFlags),
+ right = populateSpjParams(
+ smj.right, commonPartitionValues, if (!fromLeft) joinKeyPositions
else None,
+ reducers, applyPartialClustering, replicatePartitions, rightFlags))
+ case shj: ShuffledHashJoinExec =>
+ val leftFlags = nestedSpjPushFlags(joinKeyPositions, isMatchingSide =
fromLeft)
+ val rightFlags = nestedSpjPushFlags(joinKeyPositions, isMatchingSide =
!fromLeft)
+ shj.copy(
+ // Pass join-key positions only to the side that feeds this join.
+ left = populateSpjParams(
+ shj.left, commonPartitionValues, if (fromLeft) joinKeyPositions else
None,
+ reducers, applyPartialClustering, replicatePartitions, leftFlags),
+ right = populateSpjParams(
+ shj.right, commonPartitionValues, if (!fromLeft) joinKeyPositions
else None,
+ reducers, applyPartialClustering, replicatePartitions, rightFlags))
+ case other =>
+ // Child is a scan or unary operator; push directly without side
selection.
+ populateSpjParams(
+ other, commonPartitionValues, joinKeyPositions, reducers,
+ applyPartialClustering, replicatePartitions)
+ }
- private def populateJoinKeyPositions(
+ // Populate the common partition information down to the scan nodes.
+ 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,
+ flags: SpjPushFlags = SpjPushFlags()): SparkPlan = plan match {
case scan: BatchScanExec =>
+ val existing = scan.spjParams
scan.copy(
- spjParams = scan.spjParams.copy(
- joinKeyPositions = joinKeyPositions
+ spjParams = existing.copy(
+ // Overwrite common partition values only when this push targets
this side.
+ commonPartitionValues = commonPartitionValues match {
+ case Some(values) if flags.pushCommonValues => Some(values)
+ case _ => existing.commonPartitionValues
+ },
+ // Compose with positions already pushed by an inner join instead of
replacing them.
+ joinKeyPositions = if (flags.pushJoinKeyPositions) {
+ composeJoinKeyPositions(existing.joinKeyPositions,
joinKeyPositions)
+ } else {
+ existing.joinKeyPositions
+ },
+ // Keep inner reducers when the outer join does not push new ones.
+ reducers = reducers.orElse(existing.reducers),
+ applyPartialClustering = commonPartitionValues match {
+ case Some(_) if flags.pushCommonValues => applyPartialClustering
+ case _ => existing.applyPartialClustering
+ },
+ replicatePartitions = commonPartitionValues match {
+ case Some(_) if flags.pushCommonValues => replicatePartitions
+ case _ => existing.replicatePartitions
+ }
)
)
case node =>
- node.mapChildren(child => populateJoinKeyPositions(
- child, joinKeyPositions))
+ // Recurse through filters, projects, and other unary nodes above the
scan.
Review Comment:
Addressed
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/KeyGroupedPartitionedScan.scala:
##########
@@ -154,7 +196,9 @@ trait KeyGroupedPartitionedScan[T] {
// Now fill missing partition keys with empty partitions
val partitionMapping = nestGroupedPartitions.toMap
- spjParams.commonPartitionValues.get.flatMap {
+ val commonPartValues = projectCommonPartitionValuesIfNeeded(
Review Comment:
Addressed
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -2980,6 +2980,238 @@ 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(_.spjParams.joinKeyPositions.isDefined).get
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]
