ulysses-you commented on code in PR #58279:
URL: https://github.com/apache/spark/pull/58279#discussion_r3923892661
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -663,10 +666,8 @@ case class EnsureRequirements(
} else {
(unwrappedLeft, leftSpec)
}
- // Original `KeyedPartitioning` can be obtained from the child
directly if the child
- // satisfied the distribution requirement; or from the child's
child if it didn't as
- // the child must be a `GroupPartitionsExec` inserted by
`EnsureRequirement`
- // to satisfy the distribution requirement.
+ // The pre-alignment plan of the side that keeps its splits: its
partitioning
+ // still holds the original partition keys, one per input split.
Review Comment:
Fixed as suggested: `partiallyClusteredPositions` comes from the innermost
grouping, falling back to the spec's when there is none, and
`partiallyClusteredSpec` is gone. Revert-verified -- without it the updated
subset test throws exactly the measured `PartitioningCollection` invariant. The
second split for id = 1 makes the test fail on master as well, where the
stacked re-grouping duplicates rows. 83ed83560db
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -3462,6 +3462,145 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-58996: partially clustered join keeps its row count when
EnsureRequirements " +
+ "re-runs") {
+ // The storage-partitioned join branch has no shuffle of its own, so the
re-run of
+ // `EnsureRequirements` (triggered by the other branch below) reaches it.
With
+ // `numRowsPerSplit = 1` the two id = 1 rows end up in two splits, which
is what makes
+ // partial clustering replicate a side across two expected partitions.
Regrouping that
+ // replicated layout on the second pass concatenated the replicas and
replicated again,
+ // duplicating every id = 1 row.
+ val spColumns = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ createTable("sp1", spColumns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.sp1 VALUES (1, 'aa'), (1, 'ab'), (2, 'bb')")
+ createTable("sp2", spColumns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.sp2 VALUES (1, 'p'), (2, 'q')")
+
+ // Unpartitioned, so this branch's join materializes shuffle stages and is
converted to a
+ // shuffled hash join, which hands the whole plan back to
`EnsureRequirements`.
+ createTable("np1", spColumns, Array.empty)
+ sql("INSERT INTO testcat.ns.np1 VALUES (7, 'x')")
+ createTable("np2", spColumns, Array.empty)
+ sql("INSERT INTO testcat.ns.np2 VALUES (7, 'y')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true",
+ SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key ->
"100m") {
+ val df = sql(
+ """
+ |SELECT /*+ MERGE(a, b) */ a.id AS k
+ |FROM testcat.ns.sp1 a JOIN testcat.ns.sp2 b ON a.id = b.id
+ |UNION ALL
+ |SELECT c.id AS k
+ |FROM testcat.ns.np1 c JOIN testcat.ns.np2 d ON c.id = d.id
+ |""".stripMargin)
+ checkAnswer(df, Seq(Row(1L), Row(1L), Row(2L), Row(7L)))
+
+ // The re-run must leave the storage-partitioned side shuffle-free, with
the single
+ // grouping per child the first pass built: a grouping stacked over
another re-derives the
+ // alignment from an already-aligned layout and duplicates rows.
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "the storage-partitioned join must stay shuffle-free")
+ val groupPartitions =
collectGroupPartitions(df.queryExecution.executedPlan)
+ assert(groupPartitions.nonEmpty, "the storage-partitioned join must keep
its groupings")
+ groupPartitions.foreach { g =>
+ assert(collectAllGroupPartitions(g.child).isEmpty,
+ s"a GroupPartitionsExec must not be stacked over
another:\n${g.treeString}")
+ }
+ }
+ }
+
+ test("SPARK-58996: partially clustered join keeps its replicate-side choice
when " +
+ "EnsureRequirements re-runs") {
+ // The smaller side is replicated, chosen by plan statistics on the first
pass. On the re-run
+ // the statistics must be read from the pre-alignment plan again: reading
them from the
+ // aligned layout skips the statistics branch and deterministically flips
the choice. The
+ // flipped side then distributes where it used to replicate, and since the
smaller side holds
+ // more splits for id = 1 than the larger one, its raw splits overflow the
expected count
+ // (`padTo` never truncates) and the join sides end up with an unequal
number of partitions.
+ val spColumns = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ createTable("sp_small", spColumns, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.sp_small VALUES " +
+ "(1, 'a1'), (1, 'a2'), (1, 'a3'), (1, 'a4'), (1, 'a5')")
Review Comment:
Adopted -- the test now fails on master as well. One adjustment: with the
single-split shape on its own, reverting the guard stayed self-consistent in my
arms (the flipped side's one split per key fills the aligned counts), so the
data keeps the five splits for id = 1 on the smaller side too: with them the
flipped distribute side overflows the expected count of one, and reverting the
guard fails the `PartitioningCollection` invariant (measured). 83ed83560db
--
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]
