peter-toth commented on code in PR #59317:
URL: https://github.com/apache/spark/pull/59317#discussion_r4283324865
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -895,12 +895,20 @@ case class EnsureRequirements(
// nothing.
val originalKeyedPartitioning = PartitioningCollection
.representativeOf(partiallyClusteredChild.outputPartitioning).get
- val projectedOriginalPartitionKeys = partiallyClusteredPositions
- .fold(originalKeyedPartitioning.partitionKeys)(
- originalKeyedPartitioning.projectKeys(_)._2)
+ // `GroupPartitionsExec` groups the spread side as
project-then-reduce over its raw
+ // keys (see `computeGrouping`), using these same positions and
reducers. The expected
+ // count has to be computed the same way so it lands in the node's
reduced key space.
+ // Reading it 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
+ val (projectedDataTypes, projectedKeys) =
partiallyClusteredPositions.fold(
+ (originalKeyedPartitioning.keyDataTypes,
originalKeyedPartitioning.partitionKeys))(
+ originalKeyedPartitioning.projectKeys(_))
+ val reducedExpectedKeys =
spreadReducers.fold(projectedKeys)(reducers =>
Review Comment:
**Finding 2.** These lines repeat the project-then-reduce steps of
`GroupPartitionsExec.computeGrouping` (`GroupPartitionsExec.scala:716-723`).
Only the comment above keeps the two in step. The bug fixed here was exactly
this count drifting from that node's key space. A helper in `object
GroupPartitionsExec` that both call would rule out the next drift:
```scala
/** The keys a node over `childKp` groups by: projected to
`joinKeyPositions`, then reduced. */
def groupingKeys(
childKp: KeyedPartitioning,
joinKeyPositions: Option[Seq[Int]],
reducers: Option[Seq[Option[KeyReducer]]])
: (Seq[DataType], Seq[InternalRowComparableWrapper]) = {
val (projectedDataTypes, projectedKeys) =
joinKeyPositions.fold((childKp.keyDataTypes,
childKp.partitionKeys))(childKp.projectKeys)
reducers.fold((projectedDataTypes, projectedKeys))(
KeyedPartitioning.reduceKeys(projectedKeys, projectedDataTypes, _))
}
```
Here it becomes:
```scala
val reducedExpectedKeys = GroupPartitionsExec.groupingKeys(
originalKeyedPartitioning, partiallyClusteredPositions,
spreadReducers)._2
```
`computeGrouping` then starts with `val (reducedDataTypes, reducedKeys) =
groupingKeys(childKp, joinKeyPositions, reducers)`. I tried this locally (plus
an import of `DataType`). It compiles, and `GroupPartitionsExecSuite` and
`KeyGroupedPartitioningSuite` pass (260 tests).
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1328,6 +1328,32 @@ 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",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ checkAnswer(df, Seq(Row(0L), Row(1L), Row(2L), Row(3L)))
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "storage-partitioned join should be kept, not shuffled")
Review Comment:
**Finding 1.** This checks the rows and that nothing shuffles, but not that
partial clustering ran. I tried another fix that skips partial clustering
whenever either side has reducers (`&& leftReducers.isEmpty &&
rightReducers.isEmpty` on the
`v2BucketingPartiallyClusteredDistributionEnabled` check). This test still
passes with it. Both sides then just group on the reduced keys, so nothing
shuffles. The counts this PR fixes only show up in the partially clustered
layout. So I'd assert that layout too:
```scala
val plan = stripAQEPlan(df.queryExecution.executedPlan)
val gpes = collectGroupPartitions(plan)
assert(gpes.flatMap(_.ungroupingOrigin).toSet == Set(SPLIT_FOR_JOIN,
REPLICATED_FOR_JOIN),
s"one side spreads and the other repeats:\n$plan")
assert(gpes.map(_.outputPartitioning.numPartitions) == Seq(4, 4),
s"each merged key gets the two splits it holds on the spread
side:\n$plan")
val joins = collect(plan) { case smj: SortMergeJoinExec => smj }
assert(joins.size == 1 && ValidateRequirements.validate(joins.head),
s"the pair holds up:\n$plan")
```
It passes on this head. Under the skip variant it fails with `Set() did not
equal Set(SPLIT_FOR_JOIN, REPLICATED_FOR_JOIN)`. The test at
`KeyGroupedPartitioningSuite.scala:10128` checks the origins the same way.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1328,6 +1328,32 @@ 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",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
Review Comment:
**Finding 3.** This config isn't needed. Only
`KeyedShuffleSpec.canReduceKeys` reads it (`partitioning.scala:2363-2368`).
That is false whenever partial clustering is on. The pair gets in through the
`isSameFunction` arm instead, which is the point of the repro. Without this
line, the test still passes on this head and still fails on master with the
shuffle. Keeping it suggests the path depends on compatible transforms.
```suggestion
SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key
-> "true") {
```
--
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]
