Skip to content

Commit 2ced43b

Browse files
Vivek1106-04Vivek1106-04
committed
[SPARK-48290][CORE][FOLLOWUP] Bound the skewed block count and the selection cost
Address two problems with enabling the accurate recording of skewed block sizes by default. The number of accurately recorded blocks was only bounded by the size of the cutoff block, so blocks tied at that size were all recorded and could exceed spark.shuffle.maxAccurateSkewedBlockNumber by orders of magnitude. With 50000 reduce partitions, 25001 blocks of 1 KiB and 24999 blocks of 8 KiB, all 24999 larger blocks were recorded instead of at most 100. The blocks strictly larger than the cutoff are now counted up front, and the rest of the budget admits the tied blocks in block index order, which keeps the map status deterministic. Blocks at or above spark.shuffle.accurateBlockThreshold are still always recorded, as they were before. Utils.nthSmallest picked the middle element as the pivot, which is quadratic for an organ pipe distribution of block sizes. It is now an introselect: the pivot is the median of the first, middle and last element, and the range still under consideration is sorted once the ranges have stopped shrinking quickly enough, which bounds the worst case at O(numPartitions * log(numPartitions)). Both cases are covered by new tests, the ordering that the block count relies on is now asserted directly, and the benchmark also reports the organ pipe distribution.
1 parent ce414f0 commit 2ced43b

6 files changed

Lines changed: 228 additions & 53 deletions

File tree

‎core/benchmarks/HighlyCompressedMapStatusBenchmark-jdk21-results.txt‎

Lines changed: 33 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,52 @@
11
================================================================================================
2-
Build HighlyCompressedMapStatus
2+
Build HighlyCompressedMapStatus, skewed block sizes
33
================================================================================================
44

55
OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
66
Apple M4
77
2048 shuffle partitions: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
88
------------------------------------------------------------------------------------------------------------------------
9-
skewed block sizes not recorded 0 0 0 0.1 6833.0 1.0X
10-
skewed block sizes recorded accurately 0 0 0 0.1 17500.0 0.4X
9+
skewed block sizes not recorded 0 0 0 0.2 6458.0 1.0X
10+
skewed block sizes recorded accurately 0 0 0 0.1 17250.0 0.4X
1111

1212
OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
1313
Apple M4
1414
10000 shuffle partitions: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
1515
------------------------------------------------------------------------------------------------------------------------
16-
skewed block sizes not recorded 0 0 0 0.1 14041.0 1.0X
17-
skewed block sizes recorded accurately 0 0 0 0.0 68083.0 0.2X
16+
skewed block sizes not recorded 0 0 0 0.1 13834.0 1.0X
17+
skewed block sizes recorded accurately 0 0 0 0.0 56250.0 0.2X
1818

1919
OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
2020
Apple M4
2121
50000 shuffle partitions: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
2222
------------------------------------------------------------------------------------------------------------------------
23-
skewed block sizes not recorded 0 0 0 0.0 67333.0 1.0X
24-
skewed block sizes recorded accurately 0 0 0 0.0 377667.0 0.2X
23+
skewed block sizes not recorded 0 0 0 0.0 35959.0 1.0X
24+
skewed block sizes recorded accurately 0 0 0 0.0 309750.0 0.1X
25+
26+
27+
================================================================================================
28+
Build HighlyCompressedMapStatus, organ pipe block sizes
29+
================================================================================================
30+
31+
OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
32+
Apple M4
33+
2048 shuffle partitions: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
34+
------------------------------------------------------------------------------------------------------------------------
35+
skewed block sizes not recorded 0 0 0 0.2 6416.0 1.0X
36+
skewed block sizes recorded accurately 0 0 0 0.0 23583.0 0.3X
37+
38+
OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
39+
Apple M4
40+
10000 shuffle partitions: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
41+
------------------------------------------------------------------------------------------------------------------------
42+
skewed block sizes not recorded 0 0 0 0.1 14666.0 1.0X
43+
skewed block sizes recorded accurately 0 0 0 0.0 99458.0 0.1X
44+
45+
OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2
46+
Apple M4
47+
50000 shuffle partitions: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
48+
------------------------------------------------------------------------------------------------------------------------
49+
skewed block sizes not recorded 0 0 0 0.0 54583.0 1.0X
50+
skewed block sizes recorded accurately 1 1 0 0.0 530209.0 0.1X
2551

2652

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

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,11 @@ private[spark] object HighlyCompressedMapStatus {
286286
Option(SparkEnv.get)
287287
.map(_.conf.get(config.SHUFFLE_ACCURATE_BLOCK_THRESHOLD))
288288
.getOrElse(config.SHUFFLE_ACCURATE_BLOCK_THRESHOLD.defaultValue.get)
289+
// Sizes at or above `threshold` are recorded accurately. Blocks of exactly `skewCutoff` are
290+
// only recorded while `numTiedSkewedBlocksToRecord` lasts, because there can be arbitrarily
291+
// many of them and recording them all would blow past SHUFFLE_MAX_ACCURATE_SKEWED_BLOCK_NUMBER.
292+
var skewCutoff = Long.MaxValue
293+
var numTiedSkewedBlocksToRecord = 0
289294
val threshold =
290295
if (accurateBlockSkewedFactor > 0) {
291296
val maxAccurateSkewedBlockNumber =
@@ -299,10 +304,20 @@ private[spark] object HighlyCompressedMapStatus {
299304
// instead of sorting the sizes, which every map task would otherwise pay for.
300305
val sizes = uncompressedSizes.clone()
301306
val medianSize: Long = Utils.medianInPlace(sizes)
302-
val smallestAccurateSize =
303-
Utils.nthSmallest(sizes, totalNumBlocks - maxAccurateSkewedBlockNumber)
307+
val firstAccurateIdx = totalNumBlocks - maxAccurateSkewedBlockNumber
308+
skewCutoff = Utils.nthSmallest(sizes, firstAccurateIdx)
309+
// At most `maxAccurateSkewedBlockNumber` sizes are strictly larger than `skewCutoff`, and
310+
// they all sit above `firstAccurateIdx` now, so counting them bounds how many blocks of
311+
// exactly `skewCutoff` may still be recorded.
312+
var numLargerSizes = 0
313+
var j = firstAccurateIdx
314+
while (j < totalNumBlocks) {
315+
if (sizes(j) > skewCutoff) numLargerSizes += 1
316+
j += 1
317+
}
318+
numTiedSkewedBlocksToRecord = maxAccurateSkewedBlockNumber - numLargerSizes
304319
val skewSizeThreshold =
305-
Math.max(medianSize * accurateBlockSkewedFactor, smallestAccurateSize.toDouble)
320+
Math.max(medianSize * accurateBlockSkewedFactor, skewCutoff.toDouble)
306321
Math.min(shuffleAccurateBlockThreshold.toDouble, skewSizeThreshold)
307322
} else {
308323
// Disable skew detection if accurateBlockSkewedFactor <= 0
@@ -316,11 +331,21 @@ private[spark] object HighlyCompressedMapStatus {
316331
numNonEmptyBlocks += 1
317332
// Huge blocks are not included in the calculation for average size, thus size for smaller
318333
// blocks is more accurate.
319-
if (size < threshold) {
334+
var isAccurate = size >= shuffleAccurateBlockThreshold
335+
if (!isAccurate && size >= threshold) {
336+
if (size > skewCutoff) {
337+
isAccurate = true
338+
} else if (numTiedSkewedBlocksToRecord > 0) {
339+
// Ties are broken by block index so that the map status stays deterministic.
340+
numTiedSkewedBlocksToRecord -= 1
341+
isAccurate = true
342+
}
343+
}
344+
if (isAccurate) {
345+
hugeBlockSizes(i) = MapStatus.compressSize(size)
346+
} else {
320347
totalSmallBlockSize += size
321348
numSmallBlocks += 1
322-
} else {
323-
hugeBlockSizes(i) = MapStatus.compressSize(uncompressedSizes(i))
324349
}
325350
} else {
326351
emptyBlocks.add(i)

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

Lines changed: 51 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3155,44 +3155,70 @@ private[spark] object Utils
31553155
/**
31563156
* Return the n-th smallest element (0-indexed) of a long array, reordering `sizes` in place.
31573157
*
3158-
* This is a quickselect, which runs in O(sizes.length) on average. Callers that only need a few
3159-
* order statistics should prefer it over sorting the whole array, and must not rely on the
3160-
* element order of `sizes` afterwards.
3158+
* This is an introselect: a quickselect that gives up and sorts the range still under
3159+
* consideration once the ranges have stopped shrinking quickly enough, so it runs in
3160+
* O(sizes.length) on average and in O(sizes.length * log(sizes.length)) in the worst case.
3161+
* Callers that only need a few order statistics should prefer it over sorting the whole array.
3162+
* They must not rely on the element order of `sizes` afterwards, other than that no element
3163+
* before index n is larger and no element after it is smaller than the returned value.
31613164
*/
31623165
def nthSmallest(sizes: Array[Long], n: Int): Long = {
31633166
require(n >= 0 && n < sizes.length, s"n must be in [0, ${sizes.length}) but was $n")
31643167
var low = 0
31653168
var high = sizes.length - 1
3169+
// A well chosen pivot roughly halves the range, so a good run takes about log2(sizes.length)
3170+
// rounds. Inputs such as an organ pipe distribution defeat any fixed pivot choice and would
3171+
// make plain quickselect quadratic, so stop partitioning after twice that many rounds.
3172+
var remainingRounds = 2 * (32 - Integer.numberOfLeadingZeros(sizes.length))
31663173
while (low < high) {
3167-
// The middle element keeps already sorted and reverse sorted inputs, which are both common
3168-
// for shuffle block sizes, away from the quadratic worst case.
3169-
val pivot = sizes(low + (high - low) / 2)
3170-
var i = low
3171-
var j = high
3172-
while (i <= j) {
3173-
while (sizes(i) < pivot) i += 1
3174-
while (sizes(j) > pivot) j -= 1
3175-
if (i <= j) {
3176-
val tmp = sizes(i)
3177-
sizes(i) = sizes(j)
3178-
sizes(j) = tmp
3179-
i += 1
3180-
j -= 1
3181-
}
3182-
}
3183-
if (n <= j) {
3184-
high = j
3185-
} else if (n >= i) {
3186-
low = i
3187-
} else {
3188-
// The n-th element sits between the two partitions, so it is already in its final position.
3174+
if (remainingRounds == 0) {
3175+
java.util.Arrays.sort(sizes, low, high + 1)
31893176
low = n
31903177
high = n
3178+
} else {
3179+
remainingRounds -= 1
3180+
// The median of the first, middle and last element keeps already sorted and reverse
3181+
// sorted inputs, which are both common for shuffle block sizes, on the linear path.
3182+
val pivot = medianOfThree(sizes, low, low + (high - low) / 2, high)
3183+
var i = low
3184+
var j = high
3185+
while (i <= j) {
3186+
while (sizes(i) < pivot) i += 1
3187+
while (sizes(j) > pivot) j -= 1
3188+
if (i <= j) {
3189+
val tmp = sizes(i)
3190+
sizes(i) = sizes(j)
3191+
sizes(j) = tmp
3192+
i += 1
3193+
j -= 1
3194+
}
3195+
}
3196+
if (n <= j) {
3197+
high = j
3198+
} else if (n >= i) {
3199+
low = i
3200+
} else {
3201+
// The n-th element sits between the two partitions, so it is in its final position.
3202+
low = n
3203+
high = n
3204+
}
31913205
}
31923206
}
31933207
sizes(n)
31943208
}
31953209

3210+
/** Return the median of the elements of `sizes` at the three given positions. */
3211+
private def medianOfThree(sizes: Array[Long], i: Int, j: Int, k: Int): Long = {
3212+
val a = sizes(i)
3213+
val b = sizes(j)
3214+
val c = sizes(k)
3215+
if (a < b) {
3216+
if (b < c) b else if (a < c) c else a
3217+
} else {
3218+
if (a < c) a else if (b < c) c else b
3219+
}
3220+
}
3221+
31963222
/**
31973223
* Same as `median(sizes, false)`, but reorders `sizes` in place instead of sorting a copy of it.
31983224
*/

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

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ object HighlyCompressedMapStatusBenchmark extends BenchmarkBase {
4848
env
4949
}
5050

51-
private def blockSizes(numPartitions: Int): Array[Long] = {
51+
private def skewedBlockSizes(numPartitions: Int): Array[Long] = {
5252
val r = new scala.util.Random(912)
5353
Array.tabulate(numPartitions) { i =>
5454
// A few blocks are orders of magnitude larger than the rest, as in a skewed shuffle.
@@ -57,26 +57,37 @@ object HighlyCompressedMapStatusBenchmark extends BenchmarkBase {
5757
}
5858
}
5959

60+
/**
61+
* Sizes that rise and then fall. This is the worst case for selecting an order statistic with a
62+
* fixed pivot choice, so it covers the cost of the fallback to sorting.
63+
*/
64+
private def organPipeBlockSizes(numPartitions: Int): Array[Long] = {
65+
Array.tabulate(numPartitions)(i => Math.min(i, numPartitions - 1 - i).toLong)
66+
}
67+
6068
override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
6169
val envWithoutSkewedSizes = envWithSkewedFactor(-1.0)
6270
val envWithSkewedSizes = envWithSkewedFactor(5.0)
63-
runBenchmark("Build HighlyCompressedMapStatus") {
64-
Seq(2048, 10000, 50000).foreach { numPartitions =>
65-
val benchmark = new Benchmark(s"$numPartitions shuffle partitions", 1, output = output)
66-
val sizes = blockSizes(numPartitions)
71+
Seq("skewed" -> skewedBlockSizes _, "organ pipe" -> organPipeBlockSizes _).foreach {
72+
case (distribution, blockSizes) =>
73+
runBenchmark(s"Build HighlyCompressedMapStatus, $distribution block sizes") {
74+
Seq(2048, 10000, 50000).foreach { numPartitions =>
75+
val benchmark = new Benchmark(s"$numPartitions shuffle partitions", 1, output = output)
76+
val sizes = blockSizes(numPartitions)
6777

68-
benchmark.addCase("skewed block sizes not recorded", 10) { _ =>
69-
SparkEnv.set(envWithoutSkewedSizes)
70-
HighlyCompressedMapStatus(loc, sizes, 0)
71-
}
78+
benchmark.addCase("skewed block sizes not recorded", 10) { _ =>
79+
SparkEnv.set(envWithoutSkewedSizes)
80+
HighlyCompressedMapStatus(loc, sizes, 0)
81+
}
7282

73-
benchmark.addCase("skewed block sizes recorded accurately", 10) { _ =>
74-
SparkEnv.set(envWithSkewedSizes)
75-
HighlyCompressedMapStatus(loc, sizes, 0)
76-
}
83+
benchmark.addCase("skewed block sizes recorded accurately", 10) { _ =>
84+
SparkEnv.set(envWithSkewedSizes)
85+
HighlyCompressedMapStatus(loc, sizes, 0)
86+
}
7787

78-
benchmark.run()
79-
}
88+
benchmark.run()
89+
}
90+
}
8091
}
8192
SparkEnv.set(null)
8293
}

‎core/src/test/scala/org/apache/spark/scheduler/MapStatusSuite.scala‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,4 +329,54 @@ class MapStatusSuite extends SparkFunSuite {
329329
"skewed block sizes must be accurate so that AQE can detect the skew")
330330
}
331331
}
332+
333+
test("SPARK-48290: blocks tied at the cutoff size do not bypass the accurate skewed block " +
334+
"limit") {
335+
// The small blocks are the majority, so the median block size is smallBlockSize. Without the
336+
// limit, all of the tied blocks would be recorded, which for a stage of 10000 map tasks would
337+
// add more than a gigabyte of map status entries.
338+
val smallBlocksLength = 25001
339+
val tiedBlocksLength = 24999
340+
val smallBlockSize = 1024L
341+
// Far more blocks share this size than may be recorded, and it is the cutoff size itself:
342+
// it is above the median times the skew factor, so the skew threshold is exactly this size.
343+
val tiedBlockSize = 8 * 1024L
344+
val maxAccurateSkewedBlockNumber =
345+
config.SHUFFLE_MAX_ACCURATE_SKEWED_BLOCK_NUMBER.defaultValue.get
346+
347+
// No skew related config is set: this asserts the out of the box behavior.
348+
val conf = new SparkConf()
349+
val env = mock(classOf[SparkEnv])
350+
doReturn(conf).when(env).conf
351+
SparkEnv.set(env)
352+
353+
val allBlocks = createArray(smallBlocksLength, smallBlockSize) ++:
354+
createArray(tiedBlocksLength, tiedBlockSize)
355+
assert(tiedBlockSize < conf.get(config.SHUFFLE_ACCURATE_BLOCK_THRESHOLD),
356+
"the tied blocks must not be recorded as huge blocks")
357+
assert(Utils.median(allBlocks, false) *
358+
conf.get(config.SHUFFLE_ACCURATE_BLOCK_SKEWED_FACTOR) < tiedBlockSize,
359+
"the cutoff size, not the median times the skew factor, must set the skew threshold")
360+
val numSmallBlocks = allBlocks.length - maxAccurateSkewedBlockNumber
361+
val avg =
362+
(smallBlockSize * smallBlocksLength +
363+
tiedBlockSize * (tiedBlocksLength - maxAccurateSkewedBlockNumber)) / numSmallBlocks
364+
365+
val loc = BlockManagerId("a", "b", 10)
366+
val status = compressAndDecompressMapStatus(MapStatus(loc, allBlocks, 5))
367+
assert(status.isInstanceOf[HighlyCompressedMapStatus])
368+
// Ties are broken by block index, so the first maxAccurateSkewedBlockNumber tied blocks are
369+
// the ones recorded accurately.
370+
val firstAccurateBlock = smallBlocksLength
371+
for (i <- 0 until firstAccurateBlock) {
372+
assert(status.getSizeForBlock(i) === avg)
373+
}
374+
for (i <- firstAccurateBlock until firstAccurateBlock + maxAccurateSkewedBlockNumber) {
375+
assert(status.getSizeForBlock(i) === compressAndDecompressSize(tiedBlockSize))
376+
}
377+
for (i <- firstAccurateBlock + maxAccurateSkewedBlockNumber until allBlocks.length) {
378+
assert(status.getSizeForBlock(i) === avg,
379+
"no more than maxAccurateSkewedBlockNumber blocks may be recorded accurately")
380+
}
381+
}
332382
}

‎core/src/test/scala/org/apache/spark/util/UtilsSuite.scala‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import org.apache.hadoop.fs.audit.CommonAuditContext.currentAuditContext
3939
import org.apache.hadoop.ipc.{CallerContext => HadoopCallerContext}
4040
import org.apache.logging.log4j.Level
4141
import org.mockito.Mockito.when
42+
import org.scalatest.time.SpanSugar._
4243
import org.scalatestplus.mockito.MockitoSugar.mock
4344

4445
import org.apache.spark.{SparkConf, SparkException, SparkFunSuite, TaskContext}
@@ -1902,6 +1903,7 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
19021903
Array.tabulate(1000)(i => i.toLong),
19031904
Array.tabulate(1000)(i => (1000 - i).toLong),
19041905
Array.tabulate(1000)(_ => 7L),
1906+
Array.tabulate(1000)(i => math.min(i, 999 - i).toLong),
19051907
Array.fill(1000)(Random.nextLong()))
19061908
inputs.foreach { input =>
19071909
val expected = input.sorted
@@ -1911,6 +1913,41 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
19111913
}
19121914
}
19131915

1916+
test("SPARK-48290: nthSmallest is not quadratic on an organ pipe distribution") {
1917+
// Sizes that rise and then fall defeat any fixed pivot choice, so a plain quickselect needs
1918+
// more than 10^11 comparisons for these selections and does not finish in any reasonable
1919+
// time. Sorting the range still under consideration keeps each selection in the millisecond
1920+
// range, so the time limit only has to be generous enough to be free of flakiness.
1921+
val length = 1 << 20
1922+
val input = Array.tabulate(length)(i => math.min(i, length - 1 - i).toLong)
1923+
val expected = input.sorted
1924+
failAfter(60.seconds) {
1925+
Seq(0, 1, length / 2 - 1, length / 2, length - 100, length - 1).foreach { n =>
1926+
assert(Utils.nthSmallest(input.clone(), n) === expected(n))
1927+
}
1928+
}
1929+
}
1930+
1931+
test("SPARK-48290: nthSmallest partitions the array around the returned element") {
1932+
// HighlyCompressedMapStatus counts the block sizes above the cutoff by scanning only the tail
1933+
// of the array, so it relies on this ordering and not just on the returned value.
1934+
val inputs = Seq(
1935+
Array.tabulate(1000)(i => i.toLong),
1936+
Array.tabulate(1000)(i => (1000 - i).toLong),
1937+
Array.tabulate(1000)(i => math.min(i, 999 - i).toLong),
1938+
Array.tabulate(1000)(i => (i % 7).toLong),
1939+
Array.fill(1000)(Random.nextLong()))
1940+
inputs.foreach { input =>
1941+
Seq(0, 1, 17, 500, 900, 998, 999).foreach { n =>
1942+
val sizes = input.clone()
1943+
val nth = Utils.nthSmallest(sizes, n)
1944+
assert(sizes(n) === nth)
1945+
assert(sizes.take(n).forall(_ <= nth), "no element below index n may be larger")
1946+
assert(sizes.drop(n + 1).forall(_ >= nth), "no element above index n may be smaller")
1947+
}
1948+
}
1949+
}
1950+
19141951
test("SPARK-48290: nthSmallest rejects an out of range index") {
19151952
intercept[IllegalArgumentException](Utils.nthSmallest(Array(1L, 2L), -1))
19161953
intercept[IllegalArgumentException](Utils.nthSmallest(Array(1L, 2L), 2))

0 commit comments

Comments
 (0)