Make the threshold of using HighlyCompressedStatus configurable.

Author: jinxing

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

@@ -300,6 +300,13 @@ package object config {
       .booleanConf
       .createWithDefault(false)
 
+  private[spark] val SHUFFLE_HIGHLY_COMPRESSED_MAPSTATUS_THRESHOLD =
+    ConfigBuilder("spark.shuffle.highlyCompressedMapStatusThreshold")
+      .doc("Compress the size of shuffle blocks in HighlyCompressedMapStatus when the number of" +
+        "reduce partitions is above this threshold.")
+      .intConf
+      .createWithDefault(2000)
+
   private[spark] val SHUFFLE_ACCURATE_BLOCK_THRESHOLD =
     ConfigBuilder("spark.shuffle.accurateBlockThreshold")
       .doc("When we compress the size of shuffle blocks in HighlyCompressedMapStatus, we will " +

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

@@ -50,7 +50,10 @@ private[spark] sealed trait MapStatus {
 private[spark] object MapStatus {
 
   def apply(loc: BlockManagerId, uncompressedSizes: Array[Long]): MapStatus = {
-    if (uncompressedSizes.length > 2000) {
+    val threshold = Option(SparkEnv.get)
+      .map(_.conf.get(config.SHUFFLE_HIGHLY_COMPRESSED_MAPSTATUS_THRESHOLD))
+      .getOrElse(config.SHUFFLE_HIGHLY_COMPRESSED_MAPSTATUS_THRESHOLD.defaultValue.get)
+    if (uncompressedSizes.length > threshold) {
       HighlyCompressedMapStatus(loc, uncompressedSizes)
     } else {
       new CompressedMapStatus(loc, uncompressedSizes)