naveenp2708 commented on code in PR #59317:
URL: https://github.com/apache/spark/pull/59317#discussion_r4309425215
##########
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:
Thanks @peter-toth , this is the cleaner fix. Took it in 174813b:
reducersBothWays now returns no reducers when canReduceKeys is false, so the
pair stays on the raw keys, nothing merges, and the extra replicated reads go
away. I dropped the count-space change and the groupingKeys helper that went
with it, and kept the fail-fast assert as a backstop. GroupPartitionsExec,
KeyGroupedPartitioning and EnsureRequirements all pass (339 tests).
##########
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:
No longer relevant with the source fix. groupingKeys is gone along with the
count-space change, so computeGrouping is back to its original inline form.
##########
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:
Reworked the repro for the source fix. It now checks the same-function pair
is admitted without reducing (reducers.isEmpty), along with the
partial-clustering layout, correct rows and no shuffle. I confirmed it fails
with the original overflow if I revert the fix.
--
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]
