Repository navigation
[SPARK-48290][CORE] Record skewed shuffle block sizes by default so that AQE can detect skew - #57528
Vivek1106-04 wants to merge 7 commits into
Conversation
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Two default-enabled P1 regressions are reproducible at the current head: cutoff ties bypass the configured skewed-block cap, and midpoint-pivot selection becomes quadratic on valid partition-size distributions. Please see the inline comments for concrete reproductions and suggested fixes.
| val sizes = uncompressedSizes.clone() | ||
| val medianSize: Long = Utils.medianInPlace(sizes) | ||
| val smallestAccurateSize = | ||
| Utils.nthSmallest(sizes, totalNumBlocks - maxAccurateSkewedBlockNumber) |
There was a problem hiding this comment.
[P1] Enforce the skewed-block cap when the cutoff has ties
nthSmallest returns the 100th-largest block size, but the classification below records every block with size >= threshold. This means ties at the cutoff bypass spark.shuffle.maxAccurateSkewedBlockNumber.
For example, with the proposed defaults and 50,000 reduce partitions, consider 25,001 blocks of 1 KiB and 24,999 blocks of 8 KiB. The median is 1 KiB, the 100th-largest size is 8 KiB, and the threshold is 8 KiB, so all 24,999 larger blocks are recorded instead of at most 100. writeExternal writes at least five bytes per recorded entry, so a 10,000-map stage produces 1,249,950,000 bytes of uncompressed map-status entries, before the much larger driver-side collection overhead. The base branch's disabled default does not record any of these blocks.
Since this PR enables this path by default, please enforce the actual cardinality bound, including deterministic handling of ties, while preserving mandatory accurate recording above spark.shuffle.accurateBlockThreshold. Please also add a regression test containing more than 100 equal-sized blocks at the cutoff.
There was a problem hiding this comment.
You're right, and the reproduction is exact. Fixed in 2ced43b.
The cutoff size alone was never a cardinality bound, only a size bound. The blocks strictly larger than the cutoff are now counted up front -- the selection leaves them all in the top-K window, so that count is at most the configured maximum -- and the remainder of the budget admits blocks of exactly the cutoff size, in block index order, so the map status stays deterministic. No strictly larger block is ever dropped in favour of a tied one, since only ties are rationed. Blocks at or above spark.shuffle.accurateBlockThreshold are still recorded unconditionally, so the mandatory path is unchanged.
The regression test is your example verbatim: 50000 partitions, 25001 blocks of 1 KiB and 24999 of 8 KiB, stock defaults. Without the bound it fails with 1024 did not equal 4600 -- all 24999 tied blocks get recorded and avgSize collapses to the size of the small blocks, which is a neat second symptom of the same bug.
| while (low < high) { | ||
| // The middle element keeps already sorted and reverse sorted inputs, which are both common | ||
| // for shuffle block sizes, away from the quadratic worst case. | ||
| val pivot = sizes(low + (high - low) / 2) |
There was a problem hiding this comment.
[P1] Bound the worst-case cost of default-enabled quickselect
Always selecting the middle element makes the algorithm quadratic for a valid organ-pipe reducer-size distribution: Array.tabulate(n)(i => math.min(i, n - 1 - i).toLong). Replaying the exact selections made by HighlyCompressedMapStatus produces 805,003 comparisons at 2,048 partitions, 3,174,774 at 4,096, 12,682,753 at 8,192, and 50,455,932 at 16,384. Doubling the number of partitions approximately quadruples the work.
This becomes a production regression because the PR changes spark.shuffle.accurateBlockSkewedFactor from -1.0 to 5.0, making every map task in a qualifying shuffle execute these selections; the base default performs no selection. At 16,384 partitions, a 10,000-map stage would require more than 504 billion comparisons just to construct map statuses.
Please use a worst-case-bounded selection algorithm, or an introspective/randomized pivot with a safe fallback, and add an organ-pipe input to the tests and benchmark.
There was a problem hiding this comment.
Confirmed, and thanks for the organ pipe input specifically. Fixed in 2ced43b.
nthSmallest is now an introselect: median-of-three pivot, and it sorts the range still under consideration once the ranges have stopped shrinking for 2 * log2(n) rounds. That bounds the worst case at O(n log n) regardless of the input, rather than relying on the pivot choice being lucky.
One thing I want to be upfront about: I did not reproduce your comparison counts or instrument the post-fix ones. What I added is a time bound on the input you named -- six selections over a 2^20 organ pipe under an explicit 60 second limit, which currently run in 100ms total. Extrapolating your table (50,455,932 at 16384, roughly quadrupling per doubling) puts a quadratic selection at more than 10^11 comparisons there, so the limit is generous enough not to be flaky while still failing decisively on a regression. If you would rather have an actual comparison count asserted against a ceiling, I can add one, though it means an instrumented copy of the loop in test code.
I also added a test asserting that nthSmallest leaves the array partitioned around the returned element, since HighlyCompressedMapStatus now counts the sizes above the cutoff by scanning only the tail of the array and so depends on that ordering, not just on the returned value.
The benchmark now reports the organ pipe distribution alongside the skewed one. At 50000 partitions it is 530us per map task against 310us for the ordinary case -- a 1.7x gap, not a blowup.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Two independently reproduced P2 issues remain on the current head: mapper-invariant cutoff-tie selection can hide a genuinely skewed reducer, and default-enabled boxed map metadata retains approximately 11.55x the driver heap estimated in the PR. See the inline reproductions.
| if (size > skewCutoff) { | ||
| isAccurate = true | ||
| } else if (numTiedSkewedBlocksToRecord > 0) { | ||
| // Ties are broken by block index so that the map status stays deterministic. |
There was a problem hiding this comment.
[P2] Rotate cutoff ties across map tasks to avoid hiding skew
The new tie cap always admits the lowest reducer IDs, and each mapper observes the same tied reducers. MapOutputTracker.getStatistics sums each reducer's getSizeForBlock across map tasks, so a tied reducer that loses here is permanently represented by avgSize.
For example, use the defaults with 2,001 reducers and 10,000 map tasks. On every mapper, reducers 0-1899 have 10 KiB blocks and reducers 1900-2000 have 100 KiB blocks. The median is 10 KiB, the cutoff is 100 KiB, and the 100-entry cap always selects reducers 1900-1999. Reducer 2000 instead reports (1900 * 10240 + 102400) / 1901 = 10288 bytes per mapper, so AQE sees 102,880,000 bytes although its real size is 1,024,000,000 bytes. The real partition exceeds the skew threshold max(256 MiB, 5 * 102,880,000) = 514,400,000 bytes, but is never detected or split; the other 100 equally skewed reducers are detected.
Please rotate deterministic cutoff-tie selection using the mapper identity (for example mapTaskId) and add a regression test involving multiple map statuses, while preserving the 100-entry cap.
There was a problem hiding this comment.
[P2] The new rotation still leaves genuinely skewed reducers invisible
This remains reproducible at d3be2ba389d6594077219f52be0d5c5919e068a5. The new floorMod(mapTaskId, numTiedSkewedBlocks) advances a 100-entry tied-reducer window by only one position per consecutive mapper. Therefore M mappers cover only 100 + M - 1 tied reducers before wraparound.
Concrete default-configuration reproduction:
- 2,001 reduce partitions and 10 map tasks.
- Reducers 0-1880 produce 5 MiB blocks on every mapper.
- Reducers 1881-2000 produce 50 MiB blocks on every mapper, below the 100 MiB mandatory-accuracy threshold.
- The accurate-skew cap remains 100 per mapper.
Each status reports avgSize = 5,739,312 bytes for unselected blocks. Reducers 1990-2000 never enter any accurate window, so MapOutputTracker reports 57,393,120 bytes although each reducer actually contains 524,288,000 bytes. AQE computes a skew threshold of max(256 MiB, 5 * 57,393,120) = 286,965,600 bytes, so these genuinely skewed reducers remain unsplit. Overall, 19 of the 120 tied reducers miss the threshold.
There is also a second failure mode even when rotation distributes selections evenly: with 1,801 reducers producing 10 KiB blocks, 200 tied reducers producing 100 KiB blocks, and 10,000 mappers, each tied reducer is recorded on half the mappers. Unselected ties inflate avgSize to 15,087 bytes, so each skewed reducer is reported as 636,320,000 bytes against an AQE threshold of 754,350,000, even though its actual size is 1,024,000,000. All 200 genuinely skewed reducers are missed.
The new regression only covers 101 ties and 2,000 mappers, where almost every tied block is recorded on almost every mapper. Please cover both fewer-mapper/larger-tie groups and fairly sampled larger tie groups, while preserving the per-status cap and AQE visibility.
| .version("3.3.0") | ||
| .doubleConf | ||
| .createWithDefault(-1.0) | ||
| .createWithDefault(5.0) |
There was a problem hiding this comment.
[P2] Bound retained driver heap before enabling skew recording by default
Changing the default from -1.0 to 5.0 means each qualifying map status can now retain up to 100 additional entries in mutable.Map[Int, Byte], and the driver retains those deserialized map statuses. The PR estimates five bytes per entry (500 bytes per mapper; approximately 5 MB for 10,000 mappers), but those are only the writeExternal payload bytes. The live representation is a Scala mutable.HashMap with nodes, bucket arrays, and typically boxed reducer IDs.
Using the exact Scala 2.13.18 version declared by this head, I measured 10,000 mutable.HashMaps containing 100 reducer IDs in 1900-1999: they retained 57,752,240 additional heap bytes, or 5,775 bytes per mapper, versus the claimed 5,000,000 bytes - an 11.55x difference before cached serialized statuses or broadcast copies. The same distribution projects to approximately 551 MiB of additional retained heap at 100,000 mappers; the base branch's disabled default records none of these subthreshold entries.
Please measure retained MapOutputTracker heap for realistic high-map stages and use compact/primitive-backed storage or otherwise bound the driver-memory cost before enabling this path by default.
|
@sunchao! Both fixed, The tie break now rotates: the window of blocks tied at the cutoff that a map task records starts at mapTaskId mod numTied and wraps, so a tied reducer is averaged away on roughly one mapper in numTied instead of on all of them — cap unchanged, map status still deterministic, avgSize unaffected since only which blocks are recorded changes, not how many. Your example is the regression test at 2000 map statuses: it asserts each status still records exactly 100 blocks and every one of the 101 tied reducers sums to more than 98% of its real size, and without the rotation it fails at 20576000 was not greater than 2.1986692E8, your arithmetic exactly. On the heap: I reproduced your measurement first — on a real MapOutputTrackerMaster holding a 10000 map task stage, the retained map statuses are 6004 bytes per map task, 57 MiB for the stage, matching your 5775 and 57,752,240, so the five bytes in the PR body were writeExternal payload and I've corrected that claim. hugeBlockSizes is now two parallel primitive arrays sorted by reduce id that getSizeForBlock binary searches, which brings the same measurement to 668 bytes per map task and 6 MiB for the stage, about 63 MiB at 100000 mappers against your projected 551 MiB; the benchmark runs that measurement itself and prints it. Since swapping a hash lookup for a binary search invites the obvious question, I benchmarked reading every block size out of every map status too: 6.5 ns per lookup with the map versus 3.9 ns with the arrays at 50000 partitions, the reduce id no longer being boxed on every get. Two costs worth naming: the serialized statuses for that stage grow from 22 KiB to 27 KiB because each map task now records a different id set and compresses slightly worse, and the rotation does nothing when blocks strictly above the cutoff already fill the budget, since then no tie is recorded on any mapper — that is the cap's semantics rather than the tie break's, so I left it alone, but happy to revisit if you'd rather it were handled. |
7adc6da to
d3be2ba
Compare
| // The driver retains one map status per map task for the lifetime of the shuffle, so the | ||
| // accurately recorded sizes must not cost more than a few bytes each. A mutable.Map[Int, Byte] | ||
| // costs an order of magnitude more, through its nodes, bucket array and boxed reduce ids. | ||
| val smallBlocksLength = 1900 |
There was a problem hiding this comment.
[P2] Make the retained-memory regression test exercise HighlyCompressedMapStatus
The test builds exactly 1900 + 100 = 2000 partitions. However, MapStatus.apply selects HighlyCompressedMapStatus only when uncompressedSizes.length > spark.shuffle.minNumPartitionsToHighlyCompress, whose default is 2000.
Consequently, both retainedBytes(-1.0) and retainedBytes(5.0) instantiate CompressedMapStatus. That implementation ignores the skew factor, so both retained sizes are identical, perBlock == 0, and this assertion succeeds even if HighlyCompressedMapStatus regresses to the previously reported boxed-map representation.
Please use at least 2,001 partitions, for example 1,901 small blocks plus 100 skewed blocks, and assert that each measured status is a HighlyCompressedMapStatus before estimating its retained size.
There was a problem hiding this comment.
You're right, and the test was worse than under-powered: it passed because the feature never ran. 1900 + 100 = 2000 and MapStatus.apply compares with >, so both measurements built a CompressedMapStatus, perBlock was exactly 0.0, and 0.0 < 16 held. The assertion was one-sided with no floor and no type check, so "the skew path stored nothing" and "the skew path stored something compact" were indistinguishable to it.
Fixed in 4196cb4: 1,901 small blocks plus 100 skewed, an assertion that the block count exceeds spark.shuffle.minNumPartitionsToHighlyCompress, an isInstanceOf[HighlyCompressedMapStatus] check on each measured status, and a check that the expected number of skewed sizes actually came back accurate before anything is measured.
| if (size > skewCutoff) { | ||
| isAccurate = true | ||
| } else if (numTiedSkewedBlocksToRecord > 0) { | ||
| // Ties are broken by block index so that the map status stays deterministic. |
There was a problem hiding this comment.
[P2] The new rotation still leaves genuinely skewed reducers invisible
This remains reproducible at d3be2ba389d6594077219f52be0d5c5919e068a5. The new floorMod(mapTaskId, numTiedSkewedBlocks) advances a 100-entry tied-reducer window by only one position per consecutive mapper. Therefore M mappers cover only 100 + M - 1 tied reducers before wraparound.
Concrete default-configuration reproduction:
- 2,001 reduce partitions and 10 map tasks.
- Reducers 0-1880 produce 5 MiB blocks on every mapper.
- Reducers 1881-2000 produce 50 MiB blocks on every mapper, below the 100 MiB mandatory-accuracy threshold.
- The accurate-skew cap remains 100 per mapper.
Each status reports avgSize = 5,739,312 bytes for unselected blocks. Reducers 1990-2000 never enter any accurate window, so MapOutputTracker reports 57,393,120 bytes although each reducer actually contains 524,288,000 bytes. AQE computes a skew threshold of max(256 MiB, 5 * 57,393,120) = 286,965,600 bytes, so these genuinely skewed reducers remain unsplit. Overall, 19 of the 120 tied reducers miss the threshold.
There is also a second failure mode even when rotation distributes selections evenly: with 1,801 reducers producing 10 KiB blocks, 200 tied reducers producing 100 KiB blocks, and 10,000 mappers, each tied reducer is recorded on half the mappers. Unselected ties inflate avgSize to 15,087 bytes, so each skewed reducer is reported as 636,320,000 bytes against an AQE threshold of 754,350,000, even though its actual size is 1,024,000,000. All 200 genuinely skewed reducers are missed.
The new regression only covers 101 ties and 2,000 mappers, where almost every tied block is recorded on almost every mapper. Please cover both fewer-mapper/larger-tie groups and fairly sampled larger tie groups, while preserving the per-status cap and AQE visibility.
|
Both reproductions confirmed at d3be2ba, and they land on your numbers exactly. A reports a median of 57,393,120 and a threshold of 286,965,600, with 19 of the 120 tied reducers under it. B reports 636,320,000 against 754,350,000, all 200 missed. I don't think this can be fixed by rotating better, so rather than send another window schedule I want to make the case that the approach is wrong. In B the tie group is exactly twice the cap, so coverage is already a perfect 50% and stays 50% under any rotation policy, including a full-window stride. The miss doesn't come from coverage. It comes from avgSize: the 100 ties each map task drops are folded into totalSmallBlockSize, lifting avgSize to 15,088, lifting the median reducer to 150,870,000, lifting the threshold to 754,350,000. Recording more ties per map task raises the reported size of the tied reducers and the threshold they're compared against at the same time. No schedule closes that gap. A is a genuine coverage bug on top of it — a stride of one reaches only 100 + M - 1 ties — but fixing the stride alone leaves B failing. So I've stopped rationing the ties. They're all exactly skewCutoff, so the missing information is which reducers are at that size, not what their sizes are. That's one bitmap plus one byte: HighlyCompressedMapStatus now carries a run-optimized RoaringBitmap of the tied blocks alongside the existing emptyBlocks one, and the shared compressed cutoff size. Ties are excluded from the avgSize average, since they're recorded exactly and folding them in would only inflate the size reported for the blocks that really are average. The rotation, the tie counting pass and the mapTaskId-derived offset are gone; the id and size arrays still hold only the blocks strictly above the cutoff, which nthSmallest bounds. This does drop the per-status cap on tied blocks, which your P1 asked for and which you asked me to preserve here. That's deliberate — the cap is the thing that makes this unsolvable — but it's your call, so the cost is measured below rather than asserted.
Retained driver heap, SizeEstimator over the 10,000 deserialized map statuses of a 2,001 partition stage, which is what your earlier P2 asked for:
Recording all 200 ties now costs less than half of what recording 100 of them did. Two caveats. The disabled case grows 8 bytes per map task for the extra field; statuses with no ties share one never-mutated empty bitmap, which is what keeps it to that. And the bitmap's cost tracks partition count, not tie count: your 50,000-partition case from the first round, with 24,999 ties, is a single run and retains 448 bytes per status, but a maximally scattered tie set at that width would be closer to 6.25 KiB, which is worse than the capped array. Bounded by the partition count either way, which the id array was not. HighlyCompressedMapStatusBenchmark, back to back on the same machine, ns per block to build one status with sizes recorded:
Construction is flat to slightly faster — the tie counting pass is gone, which pays for the bitmap insertions — and the -1.0 cases are unchanged. getSizeForBlock gets faster, 4.7 to 3.9 ns per block at 2,048 partitions and 3.9 to 3.0 at 50,000, since the ties leave the binary-searched array. Laptop, single run per case, so treat small differences as noise and the direction as the claim. On the tests, the two rounds of misses share a cause worth naming: they asserted the mechanism — coverage, and per-reducer error against a fixed target — rather than the property. Both sides of AQE's comparison come out of the map statuses, so a test that fixes the target can't see the threshold move. The replacement computes max(skewedPartitionThresholdInBytes, skewedPartitionFactor * medianReportedSize) from the statuses themselves and requires every skewed reducer to clear it, over three shapes: your A, your B, and the 101-tie/many-mapper case. It fails on d3be2ba with 106662602 was not greater than 286965600 reducer 1881 holds 524288000 bytes but is reported as 106662602, below the AQE skew threshold of 286965600, so it would never be split One thing that fell out of writing it: the previous rotation test used 2,000 map tasks and 100 KiB tied blocks, so its tied reducers held 204,800,000 bytes — below the 256 MiB floor. They'd never be split even if reported perfectly, so that test was never measuring AQE visibility at all. The new case uses 5,000 map tasks. |
| hugeBlockIds += i | ||
| hugeBlockSizes += MapStatus.compressSize(size) | ||
| } else if (recordSkewedTies && size == skewCutoff) { | ||
| skewedBlocks.add(i) |
There was a problem hiding this comment.
[P2] Bound cutoff-tie bitmap growth for non-contiguous reducer IDs
Recording every skewCutoff tie in skewedBlocks means spark.shuffle.maxAccurateSkewedBlockNumber=100 no longer bounds retained or serialized metadata. The new regression test makes all 24,999 tied reducer IDs consecutive, so runOptimize() compresses them into one run; that is not representative of arbitrary reducer IDs.
Using this exact head and its declared RoaringBitmap 1.6.10 dependency, with 50,000 reducers, 25,001 blocks of 1 KiB, and 24,999 blocks of 8 KiB:
- Consecutive tied reducer IDs: bitmap retained payload 18 bytes; serialized bitmap 15 bytes.
- Alternating tied reducer IDs with the identical block-size distribution: retained payload 8,202 bytes; serialized bitmap 8,208 bytes, after the same
trim()andrunOptimize()calls as production.
Thus 10,000 mappers retain at least 82,020,000 bitmap bytes, and 100,000 mappers retain at least 820,200,000 bytes, despite the configured 100-entry cap. At 1,048,576 alternating reducers, the measured bitmap is 131,112 bytes per map status, or 1,311,120,000 retained bytes for 10,000 mappers. These are bitmap-payload lower bounds; object overhead and cached serialized/broadcast copies are additional, and serialized sizes above are before compression.
The current retained < 8 * 1024 test would already fail for the 50,000-partition alternating layout. Please bound tie metadata independently of reducer count while preserving AQE visibility, and add an interleaved/fragmented-cutoff-tie retained-memory and serialization regression.
|
The consecutive layout was the whole reason that bound looked comfortable: 448 bytes retained for the status, because 24,999 consecutive ids collapse into one run container. Interleaving the same size distribution retains 8,632 bytes, so the old On bounding the tie metadata independently of the reducer count, I don't think that is reachable together with what you asked for on Aug 3, and I'd rather say so than pretend otherwise. What the bitmap costs is a bit per reduce partition. That is the same bound
At the whole-status level the interleaved tie status retains 8,632 bytes, against 8,656 bytes for a status holding those very ids as empty blocks. Recording the ties costs slightly less than the bitmap this class already carried, and the new test asserts that relationship directly, so it holds across RoaringBitmap upgrades rather than resting on today's container sizes. The ties also cannot be more than half the blocks. That leaves the tension: capping tie cardinality is what makes the cost reducer-count-independent, and it is also exactly what made genuinely skewed reducers invisible in your Aug 3 reproduction, since every mapper sees the same distribution and the unrecorded ties fall back to For scale, on the two-population distribution from your own Aug 3 example (2,001 reducers, 200 tied), the bitmap is 414 bytes retained. The 8 KiB and 131 KiB figures need roughly half the reducers tied at the cutoff and scattered across every 64Ki block of the id space. An operator who does hit that and would rather not pay for it can set |
|
@sunchao! can you review the latest changes. Thank You! |
uros-b
left a comment
uros-b
left a comment
There was a problem hiding this comment.
@Vivek1106-04 Please make sure to update the migration guide accordingly (docs/core-migration-guide.md).
This flips a shipped default in a way that changes query plans (skew joins now split by default above 2000 partitions) and adds driver memory per large stage, a user-visible behavior change. Every core default flip in core-migration-guide.md carries a "Since Spark X, ... To restore the legacy behavior, set to " entry. This PR adds no such entry; the only note lives in the config .doc.
|
@uros-b added an entry to
|
|
One more minor note. Migration guide + config .doc understate worst-case driver memory. For this now-default-ON change, the docs say it retains "up to 100 block sizes per map task" but omit the new cutoff-tie RoaringBitmap, whose scattered-ties worst case is partition-proportional (~62 MiB for a 50000-partition × 10000-task stage). |
| } | ||
|
|
||
| test("SPARK-48290: nthSmallest partitions the array around the returned element") { | ||
| // HighlyCompressedMapStatus counts the block sizes above the cutoff by scanning only the tail |
There was a problem hiding this comment.
Nit: Factually wrong test comment at UtilsSuite.scala:1932-1933 claims HighlyCompressedMapStatus "scans only the tail of the array," but the shipped apply iterates the original uncompressedSizes(i) against the scalar skewCutoff and never scans the partitioned tail. Kept as a concern (not demoted) because it's a factual inaccuracy in shipped code, not just PR prose. The test itself is fine, only its stated reason is wrong.
|
I left a few more comments, @Vivek1106-04 PTAL. but otherwise lgtm! |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 4b654a06259f1d2dd4e9dc38222051d0daca39a1 against eb3b13823b1b70e131b1fb70d742cddd7b2de21d with five independent review passes. No new actionable P1/P2 findings.
The introselect fallback, cutoff-tie visibility, and compact storage changes check out. The tie bitmap still has partition-proportional memory cost; I accept that discussed design tradeoff rather than treating it as a new blocker. The existing requests to clarify the memory documentation and correct the stale test comment still apply.
Validation: all nine added regression tests passed in an isolated harness using exact-head source and cached Spark dependencies, alongside adversarial selection, map-size equivalence, serialization, and AQE-splitting probes. This was not a full Spark rebuild or an end-to-end AQE integration run. The reviewed head/base are unchanged and the current PR checks are green.
LGTM.
uros-b
left a comment
uros-b
left a comment
There was a problem hiding this comment.
Thank you @Vivek1106-04 and @sunchao!
…hat AQE can detect skew Once a shuffle has more than spark.shuffle.minNumPartitionsToHighlyCompress (2000) partitions, map statuses switch to HighlyCompressedMapStatus, which reports every block below spark.shuffle.accurateBlockThreshold as the average block size. A skewed partition whose rows are spread over many map tasks has no individual block above that threshold, so its size is averaged away, adaptive query execution sees min == median == max and the skew is left unhandled. spark.shuffle.accurateBlockSkewedFactor, added by SPARK-36967, keeps the sizes of skewed blocks accurate, but it shipped disabled because enabling it sorted the block sizes in every map task. Only two order statistics are needed, so they are now selected in O(numPartitions) instead of sorting, and the feature is enabled by default with the value its documentation already recommended.
…lection 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.
…re the accurate sizes compactly Every map task of a shuffle sees the same block size distribution, so breaking the ties at the skew cutoff by block index recorded the same reduce ids in every map status, and a reducer that lost the tie was represented by avgSize in all of them. MapOutputTracker sums those, so a genuinely skewed reducer could stay invisible to adaptive query execution. The window of tied blocks a map task records now rotates with its map task id, so a tied reducer loses on one map task out of the number of tied blocks rather than on all of them. The map status stays deterministic, avgSize is unchanged, and the per map task cap still holds. Counting the tied blocks costs one more pass over the sizes, so it is only paid when a tied block can be recorded at all. HighlyCompressedMapStatus held the accurate sizes in a mutable.Map[Int, Byte], one of which the driver retains per map task for the lifetime of the shuffle. They are now two parallel primitive arrays sorted by reduce id, which getSizeForBlock binary searches. Measured on a MapOutputTrackerMaster holding a 10000 map task stage of 2001 reducers with 101 blocks tied at the cutoff, the retained map statuses go from 6004 to 668 bytes per map task, or from 57 MiB to 6 MiB for the stage. Reading a block size also gets cheaper, from 6.5 to 3.9 nanoseconds at 50000 partitions, since the reduce id is no longer boxed on every lookup. The benchmark reports both, alongside the build cost.
…off in a bitmap Blocks tied at the skew cutoff were rationed to fit spark.shuffle.maxAccurateSkewedBlockNumber, and a rotation derived from the map task id spread which of them each map task recorded. A tie the rotation leaves out falls back to avgSize, which it raises in the process, so AQE's skew threshold rises along with the sizes it is compared against. With twice as many ties as the cap the coverage is already an even 50% under any schedule, and every tied reducer still reports below the threshold, so no window schedule fixes it. The tied blocks all share one size, so record their membership rather than their sizes: a run-optimized RoaringBitmap alongside the existing emptyBlocks one, plus the shared compressed cutoff size. Ties are excluded from the avgSize average, since they are recorded exactly. The rotation, the tie counting pass and the map task id derived offset are gone, and the id and size arrays still hold only the blocks strictly above the cutoff, which nthSmallest bounds. Statuses with no ties share one never-mutated empty bitmap. The rotation tests asserted coverage and per-reducer error against a fixed target, which cannot see the threshold move, so they are replaced by one that computes the AQE skew threshold from the map statuses themselves and requires every skewed reducer to clear it, over three tie group and map task count ratios. The retained-memory test built exactly spark.shuffle.minNumPartitionsToHighlyCompress blocks, so it measured a CompressedMapStatus and passed with the skew path never running; it now uses one more block and asserts the status type and that the sizes were recorded.
…ained memory regression The regression for recording every block tied at the skew cutoff placed all of the tied reduce ids consecutively, so the bitmap holding them compressed into a single run and retained 448 bytes. Interleaving the same size distribution produces the densest layout a bitmap can hold and retains 8,632 bytes, which the previous 8 KiB bound would have rejected. Cover both layouts, and bound them by the size of the id space rather than by a constant: a RoaringBitmap over ids below 65,536 holds one container, and a container denser than 4,096 ids becomes a fixed 8 KiB bitmap, so no layout of these ids can cost more. Assert the serialized size as well, and assert directly that recording the ties costs no more than a status holding those very ids as empty blocks, which is the bound this class has always had.
…ult in the core migration guide Flipping spark.shuffle.accurateBlockSkewedFactor from -1.0 to 5.0 changes query plans, since skewed partitions in shuffles above 2000 partitions are now visible to adaptive query execution and split by the skew join optimization, and it retains block sizes per map task on the driver. Both are user-visible, and the only note lived in the config doc. Open the "Upgrading from Core 4.3 to 4.4" section, matching the section the SQL migration guide already carries, and record the default flip along with the config that restores the legacy behavior.
…ost and fix a stale test comment The migration guide entry and the `spark.shuffle.accurateBlockSkewedFactor` documentation both described the driver memory cost as at most `spark.shuffle.maxAccurateSkewedBlockNumber` block sizes per map task, which omits the bitmap of the blocks tied at the cutoff. That bitmap is not capped: in the worst case, where the ties are scattered across the reduce id space, it costs a bit per reduce partition per map task, measured at 8202 bytes per map status for a 50000 partition stage. Both places now state that bound. The comment on the `nthSmallest` partitioning test claimed that `HighlyCompressedMapStatus` counts the sizes above the cutoff by scanning the tail of the partitioned array. That was true of an earlier revision of this change and stopped being true once the ties moved into the bitmap: `apply` compares each `uncompressedSizes(i)` against the scalar cutoff instead. The test still pins the partial ordering that `nthSmallest` documents, so only the stated reason changes.
3dc14ea to
af9cd81
Compare
…hat AQE can detect skew ### What changes were proposed in this pull request? Once a shuffle has more than `spark.shuffle.minNumPartitionsToHighlyCompress` (2000) partitions, map statuses switch to `HighlyCompressedMapStatus`, which keeps exact sizes only for blocks above `spark.shuffle.accurateBlockThreshold` (100MB) and reports the average block size for everything else. `MapOutputTracker.getStatistics` sums those per reduce id, so AQE ends up looking at a flat distribution. `spark.shuffle.accurateBlockSkewedFactor`, added by SPARK-36967, keeps skewed block sizes accurate, but it ships disabled (`-1.0`) even though its own documentation recommends setting it to the same value as `spark.sql.adaptive.skewJoin.skewedPartitionFactor`. Enabling it fully sorted the block sizes in every map task, which is presumably why it was left off. This PR: 1. replaces the sort in `HighlyCompressedMapStatus.apply` with selection of the only two order statistics it needs (the median, and the K-th largest size), using a new in-place `Utils.nthSmallest` plus `Utils.medianInPlace`. This is O(numPartitions) on average rather than O(numPartitions log numPartitions), with the same result. `nthSmallest` is an introselect: the pivot is the median of the first, middle and last element, and it sorts the range still under consideration once the ranges have stopped shrinking quickly enough, which bounds the worst case at O(numPartitions log numPartitions) even for adversarial inputs such as an organ pipe distribution; 2. enforces `spark.shuffle.maxAccurateSkewedBlockNumber` as an actual cardinality bound. Previously only the *size* of the cutoff block bounded the recording, so blocks tied at that size were all recorded. The blocks strictly larger than the cutoff are now counted up front, and the remainder of the budget admits 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; 3. changes the default of `spark.shuffle.accurateBlockSkewedFactor` from `-1.0` to `5.0`, matching `spark.sql.adaptive.skewJoin.skewedPartitionFactor`. `-1.0` still disables the feature; 4. updates the config documentation to say what leaving it off actually costs. ### Why are the changes needed? Skew join optimization silently stops working above 2000 shuffle partitions. In the reporter's case a 263GB partition spread over 10335 mappers is only ~25MB per block, far below the 100MB accurate threshold, so every block is reported as the average and AQE logs `median size == max size == min size` and finds no skew. Lowering `spark.shuffle.accurateBlockThreshold` does not help, because what matters is the per-block size rather than the per-partition total. Reproduced on master with a join where 95% of the rows land on one key, varying only `spark.sql.shuffle.partitions`: | `spark.sql.shuffle.partitions` | extra configuration | skew detected | | --- | --- | --- | | 1999 | (defaults) | yes, `SortMergeJoin(skew=true)` | | 2001 | (defaults) | no | | 4000 | (defaults) | no | | 2001 | `spark.shuffle.accurateBlockSkewedFactor=5.0` | yes | | 2001 | `spark.shuffle.minNumPartitionsToHighlyCompress=100000` | yes | Either knob restores detection on its own, which locates the problem in the map status compression rather than in AQE. The cost of the change, measured with the benchmark added here (per map task, Apple M4, JDK 21): | shuffle partitions | skewed sizes not recorded | recorded, skewed distribution | recorded, organ pipe distribution | | --- | --- | --- | --- | | 2048 | 6.4us | 17.3us | 23.6us | | 10000 | 14.7us | 56.3us | 99.5us | | 50000 | 54.6us | 309.8us | 530.2us | For comparison, the sorting implementation this PR replaces measured 28.8us, 140.7us and 1562.0us on the same machine, so at 50000 partitions the accurate path is now about 5x cheaper than sorting would have been. The organ pipe column is the worst case for the selection, where the introselect fallback kicks in; it stays within 1.7x of the ordinary case and does not grow quadratically. On the driver side each map status carries at most `spark.shuffle.maxAccurateSkewedBlockNumber` (100) extra entries of an int and a byte, so about 500 bytes per map task, or ~5MB for a 10000 map task stage. ### Does this PR introduce _any_ user-facing change? Yes. Shuffles with more than `spark.shuffle.minNumPartitionsToHighlyCompress` partitions now report skewed block sizes accurately by default, so AQE can detect and split skewed partitions where it previously could not. Queries affected by this will change plans (skew joins get split) and should run faster. Setting `spark.shuffle.accurateBlockSkewedFactor=-1.0` restores the old behavior. ### How was this patch tested? New tests in `MapStatusSuite`: - a skewed block below `spark.shuffle.accurateBlockThreshold` is recorded accurately under default configuration, and the average is not inflated by it. Fails on master with `88733 did not equal 1500`. - at 50000 blocks, 25001 of 1 KiB and 24999 of 8 KiB, exactly `spark.shuffle.maxAccurateSkewedBlockNumber` blocks are recorded and the rest report the average. Fails without the cardinality bound with `1024 did not equal 4600`, because all 24999 tied blocks were recorded and the average collapsed to the size of the small blocks. New tests in `UtilsSuite`: - `Utils.nthSmallest` agrees with a full sort over sorted, reverse sorted, constant, organ pipe, duplicate-heavy and random inputs, and its argument validation. - `Utils.nthSmallest` leaves the array partitioned around the returned element. `HighlyCompressedMapStatus` counts the sizes above the cutoff by scanning only the tail of the array, so it depends on this and not just on the returned value. - six selections over a 1048576 element organ pipe complete within a 60 second limit. A quadratic selection needs more than 10^11 comparisons there; the introselect takes 100ms in total. - `Utils.medianInPlace` agrees with `Utils.median`. `HighlyCompressedMapStatusBenchmark` is new, with results generated at 2048, 10000 and 50000 partitions for both distributions. Existing suites: `MapStatusSuite`, `MapOutputTrackerSuite`, `ShuffleSuite`, `UtilsSuite`, `org.apache.spark.shuffle.sort.*` and `AdaptiveQueryExecSuite` (129 tests) all pass. End to end, with a locally built assembly and no configuration beyond the SQL level skew join settings, the table above becomes yes / yes / yes at 1999, 2001 and 4000 partitions. I did not add an AQE level test for the >2000 partition case: `spark.shuffle.minNumPartitionsToHighlyCompress` is read through a lazy val on the JVM's `SparkEnv`, so it cannot be lowered from within a suite, and running a query with 2001 shuffle partitions takes about a minute. SPARK-36967 tested this feature at `MapStatusSuite` level for the same reason. Happy to add one if reviewers would rather pay the runtime. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code Closes #57528 from Vivek1106-04/SPARK-48290-skew-mapstatus. Authored-by: Gangavarapu Vivek <[email protected]> Signed-off-by: Chao Sun <[email protected]> (cherry picked from commit a125337) Signed-off-by: Chao Sun <[email protected]>
| Build HighlyCompressedMapStatus, skewed block sizes | ||
| ================================================================================================ | ||
|
|
||
| OpenJDK 64-Bit Server VM 21.0.10 on Mac OS X 26.5.2 |
There was a problem hiding this comment.
Hi, @Vivek1106-04 , @sunchao , @uros-b .
You should not generate the benchmark result locally.
Since this PR was generated in a wrong way, we didn't catch the Apache Spark benchmark outages.
Please see the failure.
|
I made a PR to recover the Apache Spark benchmark CI. |
…ilure when `mockito` is not on the classpath ### What changes were proposed in this pull request? This PR removes the Mockito dependency from `HighlyCompressedMapStatusBenchmark`. - Instead of mocking `SparkEnv` for each case, the benchmark now starts a single `local[1]` `SparkContext` and changes `spark.shuffle.accurateBlockSkewedFactor` on `SparkEnv.get.conf` before each case. `HighlyCompressedMapStatus` reads this value from `SparkEnv.get.conf` on every call, so the behavior is the same as the mock. - The `SparkContext` is stopped before the "Driver heap retained by MapOutputTracker" report, which creates its own `SparkContext`. ### Why are the changes needed? The benchmark job fails with: ``` java.lang.NoClassDefFoundError: org/mockito/Mockito at org.apache.spark.HighlyCompressedMapStatusBenchmark$.envWithSkewedFactor(HighlyCompressedMapStatusBenchmark.scala:52) ``` (https://github.com/dongjoon-hyun/spark/actions/runs/41684094207/job/127686502284) The GitHub Actions benchmark workflow runs `Benchmarks` through `spark-submit --jars <test jars>`. `mockito-core` is a test-scope dependency of `core`, so it is not on that classpath. `build/sbt "core/Test/runMain ..."` does include test dependencies, which is why the failure does not reproduce locally. The benchmark was added by SPARK-48290 (a125337, 2026-10-19). Other benchmarks do not depend on Mockito. - #57528 ### Does this PR introduce _any_ user-facing change? No. This only changes a benchmark. ### How was this patch tested? Ran the benchmark locally and confirmed that it completes successfully and that the reported numbers (for example, the bytes retained per map task) are the same as before the change: ``` build/sbt "core/Test/runMain org.apache.spark.HighlyCompressedMapStatusBenchmark" ``` The benchmark has not been run through `spark-submit --jars` yet, so the GitHub Actions benchmark job needs to confirm the fix. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Sonnet 5.5 Closes #60160 from dongjoon-hyun/SPARK-60984. Authored-by: Dongjoon Hyun <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>
…ilure when `mockito` is not on the classpath ### What changes were proposed in this pull request? This PR removes the Mockito dependency from `HighlyCompressedMapStatusBenchmark`. - Instead of mocking `SparkEnv` for each case, the benchmark now starts a single `local[1]` `SparkContext` and changes `spark.shuffle.accurateBlockSkewedFactor` on `SparkEnv.get.conf` before each case. `HighlyCompressedMapStatus` reads this value from `SparkEnv.get.conf` on every call, so the behavior is the same as the mock. - The `SparkContext` is stopped before the "Driver heap retained by MapOutputTracker" report, which creates its own `SparkContext`. ### Why are the changes needed? The benchmark job fails with: ``` java.lang.NoClassDefFoundError: org/mockito/Mockito at org.apache.spark.HighlyCompressedMapStatusBenchmark$.envWithSkewedFactor(HighlyCompressedMapStatusBenchmark.scala:52) ``` (https://github.com/dongjoon-hyun/spark/actions/runs/41684094207/job/127686502284) The GitHub Actions benchmark workflow runs `Benchmarks` through `spark-submit --jars <test jars>`. `mockito-core` is a test-scope dependency of `core`, so it is not on that classpath. `build/sbt "core/Test/runMain ..."` does include test dependencies, which is why the failure does not reproduce locally. The benchmark was added by SPARK-48290 (a125337, 2026-10-19). Other benchmarks do not depend on Mockito. - #57528 ### Does this PR introduce _any_ user-facing change? No. This only changes a benchmark. ### How was this patch tested? Ran the benchmark locally and confirmed that it completes successfully and that the reported numbers (for example, the bytes retained per map task) are the same as before the change: ``` build/sbt "core/Test/runMain org.apache.spark.HighlyCompressedMapStatusBenchmark" ``` The benchmark has not been run through `spark-submit --jars` yet, so the GitHub Actions benchmark job needs to confirm the fix. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Sonnet 5.5 Closes #60160 from dongjoon-hyun/SPARK-60984. Authored-by: Dongjoon Hyun <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]> (cherry picked from commit 7af0d94) Signed-off-by: Dongjoon Hyun <[email protected]>
What changes were proposed in this pull request?
Once a shuffle has more than
spark.shuffle.minNumPartitionsToHighlyCompress(2000) partitions, map statuses switch toHighlyCompressedMapStatus, which keeps exact sizes only for blocks abovespark.shuffle.accurateBlockThreshold(100MB) and reports the average block size for everything else.MapOutputTracker.getStatisticssums those per reduce id, so AQE ends up looking at a flat distribution.spark.shuffle.accurateBlockSkewedFactor, added by SPARK-36967, keeps skewed block sizes accurate, but it ships disabled (-1.0) even though its own documentation recommends setting it to the same value asspark.sql.adaptive.skewJoin.skewedPartitionFactor. Enabling it fully sorted the block sizes in every map task, which is presumably why it was left off.This PR:
HighlyCompressedMapStatus.applywith selection of the only two order statistics it needs (the median, and the K-th largest size), using a new in-placeUtils.nthSmallestplusUtils.medianInPlace. This is O(numPartitions) on average rather than O(numPartitions log numPartitions), with the same result.nthSmallestis an introselect: the pivot is the median of the first, middle and last element, and it sorts the range still under consideration once the ranges have stopped shrinking quickly enough, which bounds the worst case at O(numPartitions log numPartitions) even for adversarial inputs such as an organ pipe distribution;spark.shuffle.maxAccurateSkewedBlockNumberas an actual cardinality bound. Previously only the size of the cutoff block bounded the recording, so blocks tied at that size were all recorded. The blocks strictly larger than the cutoff are now counted up front, and the remainder of the budget admits tied blocks in block index order, which keeps the map status deterministic. Blocks at or abovespark.shuffle.accurateBlockThresholdare still always recorded, as they were before;spark.shuffle.accurateBlockSkewedFactorfrom-1.0to5.0, matchingspark.sql.adaptive.skewJoin.skewedPartitionFactor.-1.0still disables the feature;Why are the changes needed?
Skew join optimization silently stops working above 2000 shuffle partitions. In the reporter's case a 263GB partition spread over 10335 mappers is only ~25MB per block, far below the 100MB accurate threshold, so every block is reported as the average and AQE logs
median size == max size == min sizeand finds no skew. Loweringspark.shuffle.accurateBlockThresholddoes not help, because what matters is the per-block size rather than the per-partition total.Reproduced on master with a join where 95% of the rows land on one key, varying only
spark.sql.shuffle.partitions:spark.sql.shuffle.partitionsSortMergeJoin(skew=true)spark.shuffle.accurateBlockSkewedFactor=5.0spark.shuffle.minNumPartitionsToHighlyCompress=100000Either knob restores detection on its own, which locates the problem in the map status compression rather than in AQE.
The cost of the change, measured with the benchmark added here (per map task, Apple M4, JDK 21):
For comparison, the sorting implementation this PR replaces measured 28.8us, 140.7us and 1562.0us on the same machine, so at 50000 partitions the accurate path is now about 5x cheaper than sorting would have been. The organ pipe column is the worst case for the selection, where the introselect fallback kicks in; it stays within 1.7x of the ordinary case and does not grow quadratically.
On the driver side each map status carries at most
spark.shuffle.maxAccurateSkewedBlockNumber(100) extra entries of an int and a byte, so about 500 bytes per map task, or ~5MB for a 10000 map task stage.Does this PR introduce any user-facing change?
Yes. Shuffles with more than
spark.shuffle.minNumPartitionsToHighlyCompresspartitions now report skewed block sizes accurately by default, so AQE can detect and split skewed partitions where it previously could not. Queries affected by this will change plans (skew joins get split) and should run faster. Settingspark.shuffle.accurateBlockSkewedFactor=-1.0restores the old behavior.How was this patch tested?
New tests in
MapStatusSuite:spark.shuffle.accurateBlockThresholdis recorded accurately under default configuration, and the average is not inflated by it. Fails on master with88733 did not equal 1500.spark.shuffle.maxAccurateSkewedBlockNumberblocks are recorded and the rest report the average. Fails without the cardinality bound with1024 did not equal 4600, because all 24999 tied blocks were recorded and the average collapsed to the size of the small blocks.New tests in
UtilsSuite:Utils.nthSmallestagrees with a full sort over sorted, reverse sorted, constant, organ pipe, duplicate-heavy and random inputs, and its argument validation.Utils.nthSmallestleaves the array partitioned around the returned element.HighlyCompressedMapStatuscounts the sizes above the cutoff by scanning only the tail of the array, so it depends on this and not just on the returned value.Utils.medianInPlaceagrees withUtils.median.HighlyCompressedMapStatusBenchmarkis new, with results generated at 2048, 10000 and 50000 partitions for both distributions.Existing suites:
MapStatusSuite,MapOutputTrackerSuite,ShuffleSuite,UtilsSuite,org.apache.spark.shuffle.sort.*andAdaptiveQueryExecSuite(129 tests) all pass.End to end, with a locally built assembly and no configuration beyond the SQL level skew join settings, the table above becomes yes / yes / yes at 1999, 2001 and 4000 partitions.
I did not add an AQE level test for the >2000 partition case:
spark.shuffle.minNumPartitionsToHighlyCompressis read through a lazy val on the JVM'sSparkEnv, so it cannot be lowered from within a suite, and running a query with 2001 shuffle partitions takes about a minute. SPARK-36967 tested this feature atMapStatusSuitelevel for the same reason. Happy to add one if reviewers would rather pay the runtime.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code