Vivek1106-04 opened a new pull request, #57528: URL: https://github.com/apache/spark/pull/57528
### 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 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
