peter-toth commented on code in PR #58279:
URL: https://github.com/apache/spark/pull/58279#discussion_r3922990462
##########
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:
**Finding 3.** The pre-alignment plan is read here, but the positions its
keys are projected with at line 677 still come from the aligned node's spec, so
the two are in different index spaces.
`partiallyClusteredSpec` is `leftSpec`/`rightSpec`, built at line 167 from
`children(i).outputPartitioning`. On a re-run that is the reused
`GroupPartitionsExec`'s report, whose expressions are already projected
(`GroupPartitionsExec.scala:84`). So `joinKeyPositions` comes back as
`Some(Seq(0))` over one expression, while `originalKeyedPartitioning` is the
raw `[extra, id]` that `unwrapGroupPartitions` just returned.
`projectKeys(Seq(0))` then reads `extra`, `numExpectedPartitions` is keyed by
`extra` values, no merged `id` key matches, and every count stays at 1.
That is the same invariant
[r3864640493](https://github.com/apache/spark/pull/58279#discussion_r3864640493)
is about and that `applyGroupPartitions` now honours: the innermost node's
positions were computed in the raw space and stay authoritative. Both readers
of the pre-alignment plan need them, not only the writer.
[r3864640486](https://github.com/apache/spark/pull/58279#discussion_r3864640486)
named this site too — "which also fixes the stale `numExpectedPartitions`
source at lines 664-675".
**Measured.** `multi_part` partitioned by `(extra, id)` with two `extra`
values sharing `id = 1`, joined on `id`, `pushPartValues` +
`partiallyClusteredDistribution` + `allowKeysSubsetOfPartitionKeys` on, the
same `np1`/`np2` branch the other tests use:
* this head throws `java.lang.IllegalArgumentException: requirement failed:
All KeyedPartitionings in a PartitioningCollection must have equal
partitionKeys`, from `PartitioningCollection.fromPartitionings` through
`SortMergeJoinExec.outputPartitioning` inside the second pass. The distribute
side reports `[1, 1, 2]`, because `padTo(1, ...)` never truncates its two
splits, against the replicate side's `[1, 2]`.
* master returns 6 rows, `[1,10] [1,10] [1,11] [1,11] [2,20] [7,0]`, of
which 4 are correct.
The fix below is measured: the query returns the correct 4 rows, and
`KeyGroupedPartitioningSuite` + `EnsureRequirementsSuite` stay green (188).
```scala
val (partiallyClusteredChild, partiallyClusteredPositions) =
if (replicateLeftSide) {
(unwrappedRight,
innermostGroupPartition(right).flatMap(_._1.joinKeyPositions)
.orElse(rightSpec.joinKeyPositions))
} else {
(unwrappedLeft,
innermostGroupPartition(left).flatMap(_._1.joinKeyPositions)
.orElse(leftSpec.joinKeyPositions))
}
```
and at line 677:
```scala
val projectedOriginalPartitionKeys =
partiallyClusteredPositions
.fold(originalKeyedPartitioning.partitionKeys)(
originalKeyedPartitioning.projectKeys(_)._2)
```
`partiallyClusteredSpec` has no other use, so it goes away.
The subset-key test covers this shape once its left table has a second split
per join key, which also makes that test fail on master rather than only with
the `orElse` reverted:
```scala
sql("INSERT INTO testcat.ns.multi_part VALUES (1, 10, 'x'), (1, 11,
'x2'), (2, 20, 'y')")
```
with `checkAnswer(df, Seq(Row(1L, 10L), Row(1L, 11L), Row(2L, 20L), Row(7L,
0L)))`.
##########
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:
**Finding 4.** Measured: this test passes on the merge base `08757655cc2`
when run on its own, so it does not reproduce a master defect. It does pin the
guard — revert `unwrapGroupPartitions` to the one-level peel on this head and
it fails — which is the method you describe in the conversation. The comment
above it claims something stronger: "before the fix the re-run flipped the
choice and emitted an unequal number of partitions per join side". On master
this query returns the correct answer.
The reason is the data. The flip does happen on master, but `sp_small` has
no `id = 2` rows, so the `id = 2` partitions the other side over-replicates
never join anything and the extra copies stay invisible.
Putting rows on both sides of the multi-split key makes it a master-level
test. Measured: 27 rows on `08757655cc2` where 7 are correct, and it passes on
this head.
```scala
createTable("sp_small", spColumns, Array(identity("id")))
sql("INSERT INTO testcat.ns.sp_small VALUES (1, 'a'), (2, 'b')")
createTable("sp_large", spColumns, Array(identity("id")))
sql("INSERT INTO testcat.ns.sp_large VALUES " +
"(1, 'p'), (2, 'q1'), (2, 'q2'), (2, 'q3'), (2, 'q4'), (2, 'q5')")
```
with `checkAnswer(df, Row(1L) +: Seq.fill(5)(Row(2L)) :+ Row(7L))`. Left
stays the smaller side, so pass 1 still replicates it by statistics. On master
pass 2 flips, the right side's five `id = 2` partitions each hold all five of
its splits, and the join returns 25 rows for `id = 2` where 5 are correct.
--
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]
