peter-toth commented on code in PR #59317:
URL: https://github.com/apache/spark/pull/59317#discussion_r4296856585
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1328,6 +1328,46 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59234 repro: partial clustering reducer overflow") {
+ // self_reduce(id) is an identity transform that reduces itself, merging
its four keys to two
+ // (x -> x % 2) on both sides. Partial clustering spreads one side;
without the fix its expected
+ // count is read from the four un-reduced keys (one each) while the node
groups two reduced
+ // splits under each merged key, so the aligner over-emits (a planning
error with the guard, a
+ // silent shuffle on master). Reducing the expected count into the node's
space keeps the join,
+ // which this asserts: correct rows and no shuffle.
+ withFunction(SelfReduceFunction) {
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ createTable("t1", cols, Array(Expressions.apply("self_reduce",
Expressions.column("id"))))
+ sql("INSERT INTO testcat.ns.t1 VALUES (0, 'a'), (1, 'b'), (2, 'c'), (3,
'd')")
+ createTable("t2", cols, Array(Expressions.apply("self_reduce",
Expressions.column("id"))))
+ sql("INSERT INTO testcat.ns.t2 VALUES (0, 'w'), (1, 'x'), (2, 'y'), (3,
'z')")
+ val df = sql(
+ "SELECT t1.id FROM testcat.ns.t1 JOIN testcat.ns.t2 ON t1.id = t2.id")
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"true") {
+ checkAnswer(df, Seq(Row(0L), Row(1L), Row(2L), Row(3L)))
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(collectShuffles(plan).isEmpty,
+ "storage-partitioned join should be kept, not shuffled")
+ // Rows and no shuffle alone do not pin this fix: an alternative that
skips partial
+ // clustering when a side has reducers would group on the reduced
keys, shuffle nothing,
+ // and still pass. Assert the partially-clustered layout the fix
produces, so only it
+ // passes: one side spreads a merged key across tasks
(`SPLIT_FOR_JOIN`) while the other
+ // repeats it (`REPLICATED_FOR_JOIN`), each merged key carries its
four splits, and the
+ // join keeps its required distribution.
+ val gpes = collectGroupPartitions(plan)
Review Comment:
**Finding 4.** These assertions also pass a fix that never reduces. I made
`KeyedShuffleSpec.reducersBothWays` return no reducers for an `isSameFunction`
pair, on the base `EnsureRequirements`, and the test still passes:
- each side groups its four raw keys, one split each, so there is nothing to
miscount
- the nodes are still built with `SPLIT_FOR_JOIN` and `REPLICATED_FOR_JOIN`,
because partial clustering still applies
- each side still has 4 partitions
That variant never reaches the counting this PR fixes, so the comment's "so
only it passes" doesn't hold. A setup assertion closes the gap. It passes at
this head and fails on that variant:
```suggestion
val gpes = collectGroupPartitions(plan)
assert(gpes.forall(_.reducers.isDefined),
s"test setup: both sides reduce their keys:\n$plan")
```
The comment above also says "each merged key carries its four splits". Each
merged key holds two splits on the spread side, as the message of the
`numPartitions` assertion says. If you go with finding 6 instead, no side
reduces anymore, so this check would change with it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -895,12 +895,17 @@ case class EnsureRequirements(
// nothing.
val originalKeyedPartitioning = PartitioningCollection
.representativeOf(partiallyClusteredChild.outputPartitioning).get
- val projectedOriginalPartitionKeys = partiallyClusteredPositions
- .fold(originalKeyedPartitioning.partitionKeys)(
- originalKeyedPartitioning.projectKeys(_)._2)
+ // Derive the expected counts from the same keys
`GroupPartitionsExec` groups by
+ // (project then reduce, via the shared `groupingKeys`), so they
land in the node's
+ // reduced key space. Reading them from the un-reduced keys
instead lets a reducer that
+ // merges keys report a reduced key's count off an unrelated raw
key, so the aligner
+ // overflows (SPARK-59234).
+ val spreadReducers = if (replicateLeftSide) rightReducers else
leftReducers
Review Comment:
**Finding 6.** This is where a reducer reaches partial clustering. It could
be stopped one step earlier:
- Under partial clustering `KeyedShuffleSpec.canReduceKeys` is false, so
`areKeysCompatible` admits only same-function pairs
(`partitioning.scala:2306-2309`). Those need no reduce.
- `reducersBothWays` (`partitioning.scala:2388`) doesn't consult
`canReduceKeys`, though. It asks the connector anyway, and a function that
reduces itself, like `SelfReduceFunction`, returns a merging reducer.
- The `allowCompatibleTransforms` doc already says reducing keys requires
partial clustering to be off.
Making `reducersBothWays` follow the admission removes the miscount at its
source:
```scala
def reducersBothWays(other: KeyedShuffleSpec)
: (Option[Seq[Option[KeyReducer]]], Option[Seq[Option[KeyReducer]]]) =
{
// A pair `areKeysCompatible` admitted without a reduce needs no
reducers.
if (!canReduceKeys) return (None, None)
...
```
I ran it without the count change here, on the base `EnsureRequirements`.
`KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite` and
`EnsureRequirementsSuite` pass (339 tests), including your new ones. The join
then stays on the raw keys, with no merging. With 8 rows per side, both
versions give 8 partitions, but this PR adds 24 replicated partition reads and
the change above none. That number grows with how many keys the reducer merges.
The trade-off is that the count change in this PR can't be reached from SQL
anymore, though your assert still guards it. Both fixes are correct, so either
is fine with me.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -702,15 +731,9 @@ private[sql] object GroupPartitionsExec {
ungroupingOrigin: Option[UngroupingOrigin]): PartitionGrouping = {
// The routing follows the origin, the same derivation the node itself
reports.
val distributePartitions = ungroupingOrigin.exists(_ !=
REPLICATED_FOR_JOIN)
- // Project partition keys if join key positions are specified
- val (projectedDataTypes, projectedKeys) =
- joinKeyPositions.fold(
- (childKp.keyDataTypes, childKp.partitionKeys)
- )(childKp.projectKeys)
-
- // Reduce keys if reducers are specified
- val (reducedDataTypes, reducedKeys) = reducers.fold((projectedDataTypes,
projectedKeys))(
- KeyedPartitioning.reduceKeys(projectedKeys, projectedDataTypes, _))
+ // Project then reduce the child's keys. `EnsureRequirements` derives the
expected partition
+ // counts from this same call, so both land in the one reduced key space.
Review Comment:
**Finding 5.** Only the partially clustered spread side's counts come from
`groupingKeys`. The merged keys are reduced from the shuffle specs'
partitionings, and the ordering split counts the child's keys directly.
Something like:
```suggestion
// Project then reduce the child's keys. `EnsureRequirements` derives a
partially clustered
// spread side's expected counts from this same call, so both land in
one reduced key space.
```
--
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]
