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

[Vivek1106-04]

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 quickselect Utils.nthSmallest plus Utils.medianInPlace. This is O(numPartitions) rather than O(numPartitions log numPartitions), with the same result;
2. 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;
3. 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, sorting (before)	recorded, selection (this PR)
2048	6.5us	28.8us	17.5us
10000	15.0us	140.7us	68.1us
50000	67.9us	1562.0us	377.7us

At 50000 partitions the accurate path is 4.1x cheaper than it was, and the remaining overhead over not recording skewed sizes at all is 0.31ms per map task. 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 test in MapStatusSuite asserting that a skewed block below spark.shuffle.accurateBlockThreshold is recorded accurately, and that the average is not inflated by it, using default configuration. It fails on master with 88733 did not equal 1500.
• New tests in UtilsSuite checking Utils.nthSmallest against a full sort over sorted, reverse sorted, constant, duplicate-heavy and random inputs, its argument validation, and that Utils.medianInPlace agrees with Utils.median.
• New HighlyCompressedMapStatusBenchmark with results generated at 2048, 10000 and 50000 partitions.
• Existing suites: MapStatusSuite, MapOutputTrackerSuite, ShuffleSuite, UtilsSuite, org.apache.spark.shuffle.sort.* (158 tests) 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