peter-toth commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4339248460
##########
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),
Review Comment:
**Finding 6.** `reducers.orElse(existing.reducers)` has two problems:
- When the outer push brings no reducers, the inner one is kept even if the
push brings common values in another domain. With compatible transforms on,
`b12 JOIN b4` pushes a 12-to-4 reducer to `b12`, and a further join with a
second `bucket(12)` table pushes the 12 merged buckets but no reducer, so `b12`
groups its splits into 4 buckets while it reports 12.
- When the push brings reducers, they replace the inner ones instead of
composing with them, and `populateSpjParamsOnSide` passes them to both inner
sides regardless of the side flags. The reducers are indexed like the key
layout, so after the positions narrow, a kept list no longer lines up with the
projected key either.
I traced both by reading and did not run them; the three-way bucket join
above, with more than one row per bucket, would show the first. The reducers
test asserts only `reducers.nonEmpty`, so asserting the expected reducer on
each scan would pin this.
---
_Generated by [Claude Code](https://claude.ai/code)_
--
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]
