ammarchalifah commented on code in PR #59599:
URL: https://github.com/apache/spark/pull/59599#discussion_r4377510936
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -229,7 +239,8 @@ case class EnsureRequirements(
// partitioned side's BatchScanExec is grouped by join keys to
match,
// and we do that by pushing down the join keys
case Some(KeyGroupedShuffleSpec(_, _, Some(joinKeyPositions))) =>
- populateJoinKeyPositions(child, Some(joinKeyPositions))
+ populateSpjParams(
Review Comment:
Addressed
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -698,7 +716,129 @@ 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.
+ *
+ * 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)],
+ 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)
+ }
+
+ /**
+ * 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) =>
+ r.isDefined && specExpressions.exists(_.semanticEquals(e))
+ }
+ case _ => false
+ }
+ case _ => false
+ }
+ case _ => false
+ }
+ }
+
+ // Populate the common partition information down to the scan nodes,
overwriting SPJ params.
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]
