ulysses-you commented on code in PR #59863:
URL: https://github.com/apache/spark/pull/59863#discussion_r4412028521
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -924,15 +926,19 @@ case class EnsureRequirements(
// them. Each side stamps why it is left ungrouped (`UngroupingOrigin`),
so a reader of the
// finished plan judges the pair on the layouts it holds; the stamp also
decides the node's
// routing (`GroupPartitionsExec.distributePartitions`).
- (
- GroupPartitionsExec(rawLeft, leftSpec.joinKeyPositions,
- Some(mergedPartitionKeys), leftReducers,
+ def regroupOrIdentity(
+ child: SparkPlan,
+ spec: KeyedShuffleSpec,
+ reducers: Option[Seq[Option[KeyReducer]]],
+ replicates: Boolean): SparkPlan = {
+ val gpe = GroupPartitionsExec(child, spec.joinKeyPositions,
+ Some(mergedPartitionKeys), reducers,
ungroupingOrigin = Option.when(applyPartialClustering)(
- if (replicateLeftSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)),
- GroupPartitionsExec(rawRight, rightSpec.joinKeyPositions,
- Some(mergedPartitionKeys), rightReducers,
- ungroupingOrigin = Option.when(applyPartialClustering)(
- if (replicateRightSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)))
+ if (replicates) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN))
+ if (gpe.outputPartitioning == child.outputPartitioning) child else gpe
Review Comment:
Correction: a replicate-side choice can be made and then declined by the
join type. Fixed in 374a7c3f0af by keeping both nodes whenever the choice is
made, including the declined case. Added a LEFT JOIN re-run test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -924,15 +926,27 @@ case class EnsureRequirements(
// them. Each side stamps why it is left ungrouped (`UngroupingOrigin`),
so a reader of the
// finished plan judges the pair on the layouts it holds; the stamp also
decides the node's
// routing (`GroupPartitionsExec.distributePartitions`).
- (
- GroupPartitionsExec(rawLeft, leftSpec.joinKeyPositions,
- Some(mergedPartitionKeys), leftReducers,
+ def regroupOrIdentity(
+ child: SparkPlan,
+ spec: KeyedShuffleSpec,
+ reducers: Option[Seq[Option[KeyReducer]]],
+ replicates: Boolean): SparkPlan = {
+ val gpe = GroupPartitionsExec(child, spec.joinKeyPositions,
+ Some(mergedPartitionKeys), reducers,
ungroupingOrigin = Option.when(applyPartialClustering)(
- if (replicateLeftSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)),
- GroupPartitionsExec(rawRight, rightSpec.joinKeyPositions,
- Some(mergedPartitionKeys), rightReducers,
- ungroupingOrigin = Option.when(applyPartialClustering)(
- if (replicateRightSide) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN)))
+ if (replicates) REPLICATED_FOR_JOIN else SPLIT_FOR_JOIN))
+ // A spread or replicate pairing keeps its nodes even when the
grouping is the identity:
+ // a re-run peels a local sort only through a `GroupPartitionsExec`,
and that peel is how
+ // the statistics-based replicate-side choice is reproduced.
+ if (!applyPartialClustering &&
Review Comment:
Fixed in 374a7c3f0af: keep both nodes whenever a replicate side is chosen,
including when the join type declines it. Added the LEFT JOIN re-run regression
test.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -10759,26 +10891,19 @@ class KeyGroupedPartitioningSuite
val plan = df.queryExecution.executedPlan
assert(ValidateRequirements.validate(plan),
s"the executed plan must satisfy every operator's required
distribution:\n$plan")
- // Locate the marked side's node: no reducers, and output partition
i holds exactly input
- // partition i. That is weaker than `identityGrouping`, which also
asks for one output per
- // input, so the assertions below are what pin the claim and the
three keys.
- val identityGpe = collectAllGroupPartitions(plan).find { g =>
- g.reducers.isEmpty && g.groupedPartitions.zipWithIndex.forall {
- case ((_, inputIndices), outputIndex) =>
- inputIndices.lengthCompare(1) == 0 && inputIndices.head ==
outputIndex
- }
- }
- assert(identityGpe.isDefined,
- s"expected a reducer-free identity GroupPartitionsExec,
got:\n$plan")
- identityGpe.get.outputPartitioning match {
- case k: KeyedPartitioning =>
- assert(k.mayContainUnknownPartitionKeys,
- "the identity grouping must keep the unknown-keyed claim")
- assert(k.numPartitions === 3,
- s"the marked side keeps all three of its declared
keys:\n$plan")
- case other =>
- fail(s"expected a KeyedPartitioning output, got $other")
- }
+ // The marked side's alignment is the pass-through, so no node is
built for it: the
+ // claim has to reach the second join on the side itself. Assert on
that join, since the
+ // first join's keyed shuffle already carries a marked three-key
member of its own.
+ val secondJoin = collect(plan) { case j: SortMergeJoinExec => j
}.head
+ assert(collectAllShuffles(plan).size == 1,
+ s"only the first join shuffles:\n$plan")
+ val marked = collect(secondJoin.left) { case p: SparkPlan =>
p.outputPartitioning }
+ .flatMap(PartitioningCollection.representativeOf)
+ .filter(_.mayContainUnknownPartitionKeys)
Review Comment:
Fixed in 374a7c3f0af: the assertion reads the second join left input
partitioning directly.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1158,6 +1164,55 @@ class KeyGroupedPartitioningSuite
testWithCustomersAndOrders(customers_partitions, Array.empty, 2, 0)
}
+ test("SPARK-60618: Skip the pass-through GroupPartitionsExec in a
storage-partitioned " +
+ "join's alignment") {
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ // `items` reports more partition keys than the join uses, so the
aligned layout has to be
+ // projected onto the join keys. The shuffled side is laid out on that
same projection, so
+ // its `GroupPartitionsExec` would only re-state the keys its shuffle
already has and must
+ // not be built.
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true") {
+ withTable(s"testcat.ns.$items", s"testcat.ns.$purchases") {
+ createTable(items, itemsColumns, Array(identity("id"),
identity("name")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 'aa', 30.0, cast('2020-01-02' as timestamp)), " +
+ s"(3, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+ s"(4, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+ createTable(purchases, purchasesColumns, Array.empty)
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 42.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 89.0, cast('2020-01-03' as timestamp)), " +
+ s"(3, 19.5, cast('2020-02-01' as timestamp)), " +
+ s"(5, 26.0, cast('2023-01-01' as timestamp)), " +
+ s"(6, 50.0, cast('2023-02-01' as timestamp))")
+
+ val df = sql(s"SELECT i.id, i.name, i.price AS purchase_price, p.price
AS sale_price " +
+ s"FROM testcat.ns.$items i JOIN testcat.ns.$purchases p ON i.id =
p.item_id")
+ checkAnswer(df, Seq(
+ Row(1, "aa", 30.0, 42.0), Row(1, "aa", 30.0, 89.0),
+ Row(1, "aa", 40.0, 42.0), Row(1, "aa", 40.0, 89.0),
+ Row(3, "bb", 10.0, 19.5)))
+
+ val plan = df.queryExecution.executedPlan
+ // The one-side shuffle still happens; only the regrouping node over
it goes away.
+ assert(collectAllShuffles(plan).size == 1, s"expected one
shuffle:\n$plan")
+ assert(collect(plan) {
+ case g: GroupPartitionsExec
+ if g.child.isInstanceOf[ShuffleQueryStageExec] ||
+ g.child.isInstanceOf[AQEShuffleReadExec] => g
+ }.isEmpty, s"no GroupPartitionsExec should sit over a shuffle
stage:\n$plan")
+ // The source side keeps its node: `items` holds two splits for key
(1, 'aa'), and merging
+ // them is work the join needs.
+ assert(collect(plan) { case _: GroupPartitionsExec => true }.nonEmpty,
+ s"the source side should keep its grouping node:\n$plan")
Review Comment:
Fixed in 374a7c3f0af: the test enables AQE and the threshold, asserts SHJ
conversion, and checks the source-side grouping on the join input.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -1757,7 +1753,9 @@ class EnsureRequirementsSuite extends SharedSparkSession {
// The fine pair's members are each their child's own layout, so no
projection is reported and
// the nodes only align to the merged keys. Taking the first agreeing
pair instead groups both
// sides on `a` alone, i.e. positions `Some(Seq(0))`.
- assert(groupPartitionsNodes(planned).map(_.joinKeyPositions) ===
Seq(None, None),
+ // A side whose alignment is the pass-through keeps its own layout and
gets no node, so the
+ // pair shows up as unprojected nodes only.
+ assert(groupPartitionsNodes(planned).forall(_.joinKeyPositions.isEmpty),
Review Comment:
Fixed in 374a7c3f0af: both assertions now require the exact node sequence,
and the three SPARK-41413 cases each require one grouping node.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -5416,6 +5477,67 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-58996: partially clustered join keeps its replicate-side choice
when the " +
+ "alignment is the identity") {
+ // The smaller side by statistics is replicated: `pc_small` holds five
splits in 200 bytes,
+ // `pc_large` four splits in 224. Both are bucketed, whose transform
derives no ordering, so
+ // the join's sort has to be added. The larger side ends up with one split
per bucket, so the
+ // alignment on it is the identity and the pass-through node has to stay:
a re-run peels a
+ // local sort only through a `GroupPartitionsExec`, and reading the
statistics from the scans
+ // again is what keeps the choice. Without the node the counts (five
against four) pick the
+ // other side.
Review Comment:
Fixed in 374a7c3f0af: corrected the split count and size, and renamed the
test with the SPARK-60618 prefix.
--
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]
