[SPARK-48290][CORE] Record skewed shuffle block sizes by default so that AQE can detect skew

core/benchmarks/HighlyCompressedMapStatusBenchmark-jdk21-results.txt

@@ -0,0 +1,52 @@
+================================================================================================
+Build HighlyCompressedMapStatus, skewed block sizes
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
+Apple M4
+2048 shuffle partitions:                  Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
+------------------------------------------------------------------------------------------------------------------------
+skewed block sizes not recorded                       0              0           0          0.2        6458.0       1.0X
+skewed block sizes recorded accurately                0              0           0          0.1       17250.0       0.4X
+
+OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
+Apple M4
+10000 shuffle partitions:                 Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
+------------------------------------------------------------------------------------------------------------------------
+skewed block sizes not recorded                       0              0           0          0.1       13834.0       1.0X
+skewed block sizes recorded accurately                0              0           0          0.0       56250.0       0.2X
+
+OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
+Apple M4
+50000 shuffle partitions:                 Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
+------------------------------------------------------------------------------------------------------------------------
+skewed block sizes not recorded                       0              0           0          0.0       35959.0       1.0X
+skewed block sizes recorded accurately                0              0           0          0.0      309750.0       0.1X
+
+
+================================================================================================
+Build HighlyCompressedMapStatus, organ pipe block sizes
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
+Apple M4
+2048 shuffle partitions:                  Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
+------------------------------------------------------------------------------------------------------------------------
+skewed block sizes not recorded                       0              0           0          0.2        6416.0       1.0X
+skewed block sizes recorded accurately                0              0           0          0.0       23583.0       0.3X
+
+OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
+Apple M4
+10000 shuffle partitions:                 Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
+------------------------------------------------------------------------------------------------------------------------
+skewed block sizes not recorded                       0              0           0          0.1       14666.0       1.0X
+skewed block sizes recorded accurately                0              0           0          0.0       99458.0       0.1X
+
+OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
+Apple M4
+50000 shuffle partitions:                 Best Time(ms)   Avg Time(ms)   Stdev(ms)    Rate(M/s)   Per Row(ns)   Relative
+------------------------------------------------------------------------------------------------------------------------
+skewed block sizes not recorded                       0              0           0          0.0       54583.0       1.0X
+skewed block sizes recorded accurately                1              1           0          0.0      530209.0       0.1X
+
+

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -1441,12 +1441,16 @@ package object config {
     ConfigBuilder("spark.shuffle.accurateBlockSkewedFactor")
       .doc("A shuffle block is considered as skewed and will be accurately recorded in " +
         "HighlyCompressedMapStatus if its size is larger than this factor multiplying " +
-        "the median shuffle block size or SHUFFLE_ACCURATE_BLOCK_THRESHOLD. It is " +
-        "recommended to set this parameter to be the same as SKEW_JOIN_SKEWED_PARTITION_FACTOR." +
-        "Set to -1.0 to disable this feature by default.")
+        "the median shuffle block size or SHUFFLE_ACCURATE_BLOCK_THRESHOLD. Otherwise every " +
+        "block below SHUFFLE_ACCURATE_BLOCK_THRESHOLD is reported as the average block size, " +
+        "which hides skew from adaptive query execution once a shuffle has more than " +
+        "SHUFFLE_MIN_NUM_PARTS_TO_HIGHLY_COMPRESS partitions. At most " +
+        "SHUFFLE_MAX_ACCURATE_SKEWED_BLOCK_NUMBER block sizes are recorded accurately per map " +
+        "task. The default matches SKEW_JOIN_SKEWED_PARTITION_FACTOR. Set to -1.0 to disable " +
+        "this feature.")
       .version("3.3.0")
       .doubleConf
-      .createWithDefault(-1.0)
+      .createWithDefault(5.0)
 
   private[spark] val SHUFFLE_MAX_ACCURATE_SKEWED_BLOCK_NUMBER =
     ConfigBuilder("spark.shuffle.maxAccurateSkewedBlockNumber")

core/src/main/scala/org/apache/spark/scheduler/MapStatus.scala

@@ -286,22 +286,38 @@ private[spark] object HighlyCompressedMapStatus {
       Option(SparkEnv.get)
         .map(_.conf.get(config.SHUFFLE_ACCURATE_BLOCK_THRESHOLD))
         .getOrElse(config.SHUFFLE_ACCURATE_BLOCK_THRESHOLD.defaultValue.get)
+    // Sizes at or above `threshold` are recorded accurately. Blocks of exactly `skewCutoff` are
+    // only recorded while `numTiedSkewedBlocksToRecord` lasts, because there can be arbitrarily
+    // many of them and recording them all would blow past SHUFFLE_MAX_ACCURATE_SKEWED_BLOCK_NUMBER.
+    var skewCutoff = Long.MaxValue
+    var numTiedSkewedBlocksToRecord = 0
     val threshold =
       if (accurateBlockSkewedFactor > 0) {
-        val sortedSizes = uncompressedSizes.sorted
-        val medianSize: Long = Utils.median(sortedSizes, true)
         val maxAccurateSkewedBlockNumber =
           Math.min(
             Option(SparkEnv.get)
               .map(_.conf.get(config.SHUFFLE_MAX_ACCURATE_SKEWED_BLOCK_NUMBER))
               .getOrElse(config.SHUFFLE_MAX_ACCURATE_SKEWED_BLOCK_NUMBER.defaultValue.get),
             totalNumBlocks
           )
+        // Only two order statistics are needed here, so they are selected in O(totalNumBlocks)
+        // instead of sorting the sizes, which every map task would otherwise pay for.
+        val sizes = uncompressedSizes.clone()
+        val medianSize: Long = Utils.medianInPlace(sizes)
+        val firstAccurateIdx = totalNumBlocks - maxAccurateSkewedBlockNumber
+        skewCutoff = Utils.nthSmallest(sizes, firstAccurateIdx)
+        // At most `maxAccurateSkewedBlockNumber` sizes are strictly larger than `skewCutoff`, and
+        // they all sit above `firstAccurateIdx` now, so counting them bounds how many blocks of
+        // exactly `skewCutoff` may still be recorded.
+        var numLargerSizes = 0
+        var j = firstAccurateIdx
+        while (j < totalNumBlocks) {
+          if (sizes(j) > skewCutoff) numLargerSizes += 1
+          j += 1
+        }
+        numTiedSkewedBlocksToRecord = maxAccurateSkewedBlockNumber - numLargerSizes
         val skewSizeThreshold =
-          Math.max(
-            medianSize * accurateBlockSkewedFactor,
-            sortedSizes(totalNumBlocks - maxAccurateSkewedBlockNumber).toDouble
-          )
+          Math.max(medianSize * accurateBlockSkewedFactor, skewCutoff.toDouble)
         Math.min(shuffleAccurateBlockThreshold.toDouble, skewSizeThreshold)
       } else {
         // Disable skew detection if accurateBlockSkewedFactor <= 0
@@ -315,11 +331,21 @@ private[spark] object HighlyCompressedMapStatus {
         numNonEmptyBlocks += 1
         // Huge blocks are not included in the calculation for average size, thus size for smaller
         // blocks is more accurate.
-        if (size < threshold) {
+        var isAccurate = size >= shuffleAccurateBlockThreshold
+        if (!isAccurate && size >= threshold) {
+          if (size > skewCutoff) {
+            isAccurate = true
+          } else if (numTiedSkewedBlocksToRecord > 0) {
+            // Ties are broken by block index so that the map status stays deterministic.
+            numTiedSkewedBlocksToRecord -= 1
+            isAccurate = true
+          }
+        }
+        if (isAccurate) {
+          hugeBlockSizes(i) = MapStatus.compressSize(size)
+        } else {
           totalSmallBlockSize += size
           numSmallBlocks += 1
-        } else {
-          hugeBlockSizes(i) = MapStatus.compressSize(uncompressedSizes(i))
         }
       } else {
         emptyBlocks.add(i)

core/src/main/scala/org/apache/spark/util/Utils.scala

@@ -3152,6 +3152,85 @@ private[spark] object Utils
     }
   }
 
+  /**
+   * Return the n-th smallest element (0-indexed) of a long array, reordering `sizes` in place.
+   *
+   * This is an introselect: a quickselect that gives up and sorts the range still under
+   * consideration once the ranges have stopped shrinking quickly enough, so it runs in
+   * O(sizes.length) on average and in O(sizes.length * log(sizes.length)) in the worst case.
+   * Callers that only need a few order statistics should prefer it over sorting the whole array.
+   * They must not rely on the element order of `sizes` afterwards, other than that no element
+   * before index n is larger and no element after it is smaller than the returned value.
+   */
+  def nthSmallest(sizes: Array[Long], n: Int): Long = {
+    require(n >= 0 && n < sizes.length, s"n must be in [0, ${sizes.length}) but was $n")
+    var low = 0
+    var high = sizes.length - 1
+    // A well chosen pivot roughly halves the range, so a good run takes about log2(sizes.length)
+    // rounds. Inputs such as an organ pipe distribution defeat any fixed pivot choice and would
+    // make plain quickselect quadratic, so stop partitioning after twice that many rounds.
+    var remainingRounds = 2 * (32 - Integer.numberOfLeadingZeros(sizes.length))
+    while (low < high) {
+      if (remainingRounds == 0) {
+        java.util.Arrays.sort(sizes, low, high + 1)
+        low = n
+        high = n
+      } else {
+        remainingRounds -= 1
+        // The median of the first, middle and last element keeps already sorted and reverse
+        // sorted inputs, which are both common for shuffle block sizes, on the linear path.
+        val pivot = medianOfThree(sizes, low, low + (high - low) / 2, high)
+        var i = low
+        var j = high
+        while (i <= j) {
+          while (sizes(i) < pivot) i += 1
+          while (sizes(j) > pivot) j -= 1
+          if (i <= j) {
+            val tmp = sizes(i)
+            sizes(i) = sizes(j)
+            sizes(j) = tmp
+            i += 1
+            j -= 1
+          }
+        }
+        if (n <= j) {
+          high = j
+        } else if (n >= i) {
+          low = i
+        } else {
+          // The n-th element sits between the two partitions, so it is in its final position.
+          low = n
+          high = n
+        }
+      }
+    }
+    sizes(n)
+  }
+
+  /** Return the median of the elements of `sizes` at the three given positions. */
+  private def medianOfThree(sizes: Array[Long], i: Int, j: Int, k: Int): Long = {
+    val a = sizes(i)
+    val b = sizes(j)
+    val c = sizes(k)
+    if (a < b) {
+      if (b < c) b else if (a < c) c else a
+    } else {
+      if (a < c) a else if (b < c) c else b
+    }
+  }
+
+  /**
+   * Same as `median(sizes, false)`, but reorders `sizes` in place instead of sorting a copy of it.
+   */
+  def medianInPlace(sizes: Array[Long]): Long = {
+    val len = sizes.length
+    len match {
+      case _ if (len % 2 == 0) =>
+        math.max((nthSmallest(sizes, len / 2) + nthSmallest(sizes, len / 2 - 1)) / 2, 1)
+      case _ => math.max(nthSmallest(sizes, len / 2), 1)
+    }
+  }
+
   /**
    * Check if a command is available.
    */

core/src/test/scala/org/apache/spark/HighlyCompressedMapStatusBenchmark.scala

@@ -0,0 +1,94 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark
+
+import org.mockito.Mockito.{doReturn, mock}
+
+import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
+import org.apache.spark.internal.config
+import org.apache.spark.scheduler.HighlyCompressedMapStatus
+import org.apache.spark.storage.BlockManagerId
+
+/**
+ * Benchmark for building a HighlyCompressedMapStatus, which every map task of a shuffle with more
+ * than spark.shuffle.minNumPartitionsToHighlyCompress partitions has to pay for. It covers both
+ * the case where skewed block sizes are recorded accurately and the case where they are not.
+ * {{{
+ *   To run this benchmark:
+ *   1. without sbt: bin/spark-submit --class <this class> <spark core test jar>
+ *   2. build/sbt "core/Test/runMain <this class>"
+ *   3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt "core/Test/runMain <this class>"
+ *      Results will be written to "benchmarks/HighlyCompressedMapStatusBenchmark-results.txt".
+ * }}}
+ */
+object HighlyCompressedMapStatusBenchmark extends BenchmarkBase {
+
+  private val loc = BlockManagerId("a", "host", 1000)
+
+  private def envWithSkewedFactor(skewedFactor: Double): SparkEnv = {
+    val conf = new SparkConf()
+      .set(config.SHUFFLE_ACCURATE_BLOCK_SKEWED_FACTOR.key, skewedFactor.toString)
+    val env = mock(classOf[SparkEnv])
+    doReturn(conf, Seq.empty: _*).when(env).conf
+    env
+  }
+
+  private def skewedBlockSizes(numPartitions: Int): Array[Long] = {
+    val r = new scala.util.Random(912)
+    Array.tabulate(numPartitions) { i =>
+      // A few blocks are orders of magnitude larger than the rest, as in a skewed shuffle.
+      if (i % 500 == 0) (r.nextDouble() * 512 * 1024 * 1024).toLong
+      else (r.nextDouble() * 64 * 1024).toLong
+    }
+  }
+
+  /**
+   * Sizes that rise and then fall. This is the worst case for selecting an order statistic with a
+   * fixed pivot choice, so it covers the cost of the fallback to sorting.
+   */
+  private def organPipeBlockSizes(numPartitions: Int): Array[Long] = {
+    Array.tabulate(numPartitions)(i => Math.min(i, numPartitions - 1 - i).toLong)
+  }
+
+  override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+    val envWithoutSkewedSizes = envWithSkewedFactor(-1.0)
+    val envWithSkewedSizes = envWithSkewedFactor(5.0)
+    Seq("skewed" -> skewedBlockSizes _, "organ pipe" -> organPipeBlockSizes _).foreach {
+      case (distribution, blockSizes) =>
+        runBenchmark(s"Build HighlyCompressedMapStatus, $distribution block sizes") {
+          Seq(2048, 10000, 50000).foreach { numPartitions =>
+            val benchmark = new Benchmark(s"$numPartitions shuffle partitions", 1, output = output)
+            val sizes = blockSizes(numPartitions)
+
+            benchmark.addCase("skewed block sizes not recorded", 10) { _ =>
+              SparkEnv.set(envWithoutSkewedSizes)
+              HighlyCompressedMapStatus(loc, sizes, 0)
+            }
+
+            benchmark.addCase("skewed block sizes recorded accurately", 10) { _ =>
+              SparkEnv.set(envWithSkewedSizes)
+              HighlyCompressedMapStatus(loc, sizes, 0)
+            }
+
+            benchmark.run()
+          }
+        }
+    }
+    SparkEnv.set(null)
+  }
+}