(spark) branch master updated: [SPARK-58290][SQL] Fix O(n^2) partition index access in ShufflePartitionsUtil.coalescePartitionsWithSkew

[we...@apache.org]

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 <ul...@gmail.com>
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 <ul...@gmail.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../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: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org