This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 3082cdfda92f [SPARK-58290][SQL] Fix O(n^2) partition index access in
ShufflePartitionsUtil.coalescePartitionsWithSkew
3082cdfda92f is described below
commit 3082cdfda92f4507542ccbf99cf777a51f29414c
Author: Xiduo You <[email protected]>
AuthorDate: Thu Jul 23 21:05:21 2026 +0800
[SPARK-58290][SQL] Fix O(n^2) partition index access in
ShufflePartitionsUtil.coalescePartitionsWithSkew
### What changes were proposed in this pull request?
In `ShufflePartitionsUtil.coalescePartitionsWithSkew`, the local
`partitionIndices` is materialized from `inputPartitionSpecs`, whose inner
collection is a `List` in practice. The coalescing loop then repeatedly
accesses this sequence by index (`partitionIndices(i - 1)`,
`partitionIndices(i)`, and the inner skew-section scan
`partitionIndices(repeatIndex)`). Because `List.apply(i)` is O(i), the loop
degrades to O(n^2) in the number of shuffle partitions.
This PR calls `.toArray` on `partitionIndices` so all indexed accesses
become O(1), restoring the loop to O(n). The change is local; all downstream
usages (`.head`, `.last`, indexed access, `.length`) are valid on `Array`, and
the sequence is only read.
### Why are the changes needed?
A production job hung, and the driver thread dump showed it stuck in
`ShufflePartitionsUtil.coalescePartitionsWithSkew`, spinning inside `List.drop`
/ `List.apply` (indexed list access):
```
scala.collection.immutable.List.drop(List.scala:79)
scala.collection.LinearSeqOps.apply(LinearSeq.scala:130)
scala.collection.immutable.List.apply(List.scala:79)
org.apache.spark.sql.execution.adaptive.ShufflePartitionsUtil$.coalescePartitionsWithSkew(ShufflePartitionsUtil.scala:166)
org.apache.spark.sql.execution.adaptive.ShufflePartitionsUtil$.coalescePartitions(ShufflePartitionsUtil.scala:75)
```
With a large number of shuffle partitions (skew join scenarios often
produce many), the O(n^2) behavior can hang the driver during AQE planning.
Local micro-benchmark driving `coalescePartitions` on the skew path, `List`
vs `Array` (same machine, same inputs):
| partitions | List (ms) | Array (ms) |
|-----------:|----------:|-----------:|
| 5000 | 72 | <1 |
| 10000 | 291 | <1 |
| 20000 | 1166 | 1 |
The `List` version's per-element cost grows linearly with partition count
(confirming O(n^2)), while the `Array` version stays flat (O(n)).
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests, all passing locally (39 tests):
```
build/sbt "sql/testOnly *ShufflePartitionsUtilSuite
*CoalesceShufflePartitionsSuite"
```
This is a pure performance fix with no behavior change.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Closes #57459 from ulysses-you/SPARK-58290.
Authored-by: Xiduo You <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../spark/sql/execution/adaptive/ShufflePartitionsUtil.scala | 7 +++++--
1 file changed, 5 insertions(+), 2 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
index e0bad017b8f8..b8fdaa6154e7 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
@@ -159,8 +159,11 @@ object ShufflePartitionsUtil extends Logging {
}
// The indices may look like [0, 1, 2, 2, 2, 3, 4, 4, 5], and the repeated
`2` and `4` mean
- // skewed partitions.
- val partitionIndices = partitionIndicesSeq.head
+ // skewed partitions. Materialize into an `Array` for O(1) indexed access:
the loop below
+ // indexes into this sequence repeatedly, which degrades to O(n^2) when it
is backed by a
+ // `List` (the default here) since `List.apply(i)` is O(i). With a large
number of shuffle
+ // partitions this quadratic behavior can hang the driver.
+ val partitionIndices = partitionIndicesSeq.head.toArray
// The fist index must be 0.
assert(partitionIndices.head == 0)
val newPartitionSpecsSeq = Seq.fill(mapOutputStatistics.length)(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
