Enable/disable caching of partition lengths with a configuration variable.

Signed-off-by: Pascal Spörri <psp@zurich.ibm.com>

Author: pspoerri

README.md

@@ -38,6 +38,8 @@ These configuration values need to be passed to Spark to load and configure the
 
 Changing these values might have an impact on performance.
 
+- `spark.shuffle.s3.cachePartitionLengths`: Cache partition lengths in memory (default: `true`)
+- `spark.shuffle.s3.cacheChecksums`: Cache checksums in memory (default: `true`)
 - `spark.shuffle.s3.cleanup`: Cleanup the shuffle files (default: `true`)
 - `spark.shuffle.s3.folderPrefixes`: The number of prefixes to use when storing files on S3
   (default: `10`, minimum: `1`).

src/main/scala/org/apache/spark/shuffle/helper/S3ShuffleDispatcher.scala

@@ -33,6 +33,8 @@ class S3ShuffleDispatcher extends Logging {
   private val isS3A = rootDir.startsWith("s3a://")
 
   // Optional
+  val cachePartitionLengths: Boolean = conf.getBoolean("spark.shuffle.s3.cachePartitionLengths", defaultValue = true)
+  val cacheChecksums: Boolean = conf.getBoolean("spark.shuffle.s3.cacheChecksums", defaultValue = true)
   val cleanupShuffleFiles: Boolean = conf.getBoolean("spark.shuffle.s3.cleanup", defaultValue = true)
   val folderPrefixes: Int = conf.getInt("spark.shuffle.s3.folderPrefixes", defaultValue = 10)
   val prefetchBatchSize: Int = conf.getInt("spark.shuffle.s3.prefetchBatchSize", defaultValue = 25)
@@ -57,6 +59,8 @@ class S3ShuffleDispatcher extends Logging {
   logInfo(s"- spark.shuffle.s3.rootDir=${rootDir} (app dir: ${appDir})")
 
   // Optional
+  logInfo(s"- spark.shuffle.s3.cachePartitionLengths=${cachePartitionLengths}")
+  logInfo(s"- spark.shuffle.s3.cacheChecksums=${cacheChecksums}")
   logInfo(s"- spark.shuffle.s3.cleanup=${cleanupShuffleFiles}")
   logInfo(s"- spark.shuffle.s3.folderPrefixes=${folderPrefixes}")
   logInfo(s"- spark.shuffle.s3.prefetchBlockSize=${prefetchBatchSize}")

src/main/scala/org/apache/spark/shuffle/helper/S3ShuffleHelper.scala

@@ -27,9 +27,15 @@ object S3ShuffleHelper extends Logging {
    *
    * @param shuffleIndex
    */
-  def purgeCachedShuffleIndices(shuffleIndex: Int): Unit = {
-    val blockFilter = (block: ShuffleIndexBlockId) => block.shuffleId == shuffleIndex
-    cachedArrayLengths.remove(blockFilter, None)
+  def purgeCachedDataForShuffle(shuffleIndex: Int): Unit = {
+    if (dispatcher.cachePartitionLengths) {
+      val filter = (block: ShuffleIndexBlockId) => block.shuffleId == shuffleIndex
+      cachedArrayLengths.remove(filter, None)
+    }
+    if (dispatcher.cacheChecksums) {
+      val filter = (block: ShuffleChecksumBlockId) => block.shuffleId == shuffleIndex
+      cachedChecksums.remove(filter, None)
+    }
   }
 
   /**
@@ -59,6 +65,7 @@ object S3ShuffleHelper extends Logging {
   def listShuffleIndices(shuffleId: Int): Array[ShuffleIndexBlockId] = {
     val shuffleIndexFilter: PathFilter = new PathFilter() {
       private val prefix = f"shuffle_${shuffleId}_"
+
       override def accept(path: Path): Boolean = {
         val name = path.getName
         name.startsWith(prefix) && name.endsWith("_0.index")
@@ -85,8 +92,8 @@ object S3ShuffleHelper extends Logging {
    * @param mapId
    * @return
    */
-  def getPartitionLengthsCached(shuffleId: Int, mapId: Long): Array[Long] = {
-    getPartitionLengthsCached(ShuffleIndexBlockId(shuffleId, mapId, NOOP_REDUCE_ID))
+  def getPartitionLengths(shuffleId: Int, mapId: Long): Array[Long] = {
+    getPartitionLengths(ShuffleIndexBlockId(shuffleId, mapId, NOOP_REDUCE_ID))
   }
 
   /**
@@ -95,19 +102,21 @@ object S3ShuffleHelper extends Logging {
    * @param blockId
    * @return
    */
-  def getPartitionLengthsCached(blockId: ShuffleIndexBlockId): Array[Long] = {
-    cachedArrayLengths.getOrElsePut(blockId, readBlockAsArray)
-  }
-
-  def getChecksumsCached(shuffleId: Int, mapId: Long): Array[Long] = {
-    cachedChecksums.getOrElsePut(ShuffleChecksumBlockId(shuffleId, mapId, 0), readBlockAsArray)
+  def getPartitionLengths(blockId: ShuffleIndexBlockId): Array[Long] = {
+    if (dispatcher.cachePartitionLengths) {
+      return cachedArrayLengths.getOrElsePut(blockId, readBlockAsArray)
+    }
+    readBlockAsArray(blockId)
   }
 
   def getChecksums(shuffleId: Int, mapId: Long): Array[Long] = {
-    getChecksums(ShuffleChecksumBlockId(shuffleId = shuffleId, mapId = mapId, reduceId = 0))
+    getChecksums(ShuffleChecksumBlockId(shuffleId, mapId, 0))
   }
 
   def getChecksums(blockId: ShuffleChecksumBlockId): Array[Long] = {
+    if (dispatcher.cacheChecksums) {
+      return cachedChecksums.getOrElsePut(blockId, readBlockAsArray)
+    }
     readBlockAsArray(blockId)
   }
 

src/main/scala/org/apache/spark/shuffle/sort/S3ShuffleManager.scala

@@ -133,7 +133,7 @@ private[spark] class S3ShuffleManager(conf: SparkConf) extends ShuffleManager wi
     // Remove and close all input streams.
     dispatcher.closeCachedBlocks(shuffleId)
     // Remove metadata.
-    S3ShuffleHelper.purgeCachedShuffleIndices(shuffleId)
+    S3ShuffleHelper.purgeCachedDataForShuffle(shuffleId)
   }
 
   /** Remove a shuffle's metadata from the ShuffleManager. */

src/main/scala/org/apache/spark/storage/S3ChecksumValidationStream.scala

@@ -27,12 +27,13 @@ class S3ChecksumValidationStream(
   }
 
   private val checksum: Checksum = S3ShuffleHelper.createChecksumAlgorithm(checksumAlgorithm)
-  private val lengths: Array[Long] = S3ShuffleHelper.getPartitionLengthsCached(shuffleId, mapId)
-  private val referenceChecksums: Array[Long] = S3ShuffleHelper.getChecksumsCached(shuffleId, mapId)
+  private val lengths: Array[Long] = S3ShuffleHelper.getPartitionLengths(shuffleId, mapId)
+  private val referenceChecksums: Array[Long] = S3ShuffleHelper.getChecksums(shuffleId, mapId)
 
   private var pos: Long = 0
   private var reduceId: Int = startReduceId
   private var blockLength: Long = lengths(reduceId)
+
   private def eof(): Boolean = reduceId > endReduceId
 
   validateChecksum()

src/main/scala/org/apache/spark/storage/S3ShuffleBlockIterator.scala

@@ -57,7 +57,7 @@ class S3ShuffleBlockIterator(
   }
 
   def getAccumulatedLengths(shuffleId: Int, mapId: Long): Array[Long] = {
-    val lengths = S3ShuffleHelper.getPartitionLengthsCached(shuffleId, mapId)
+    val lengths = S3ShuffleHelper.getPartitionLengths(shuffleId, mapId)
     Array[Long](0) ++ lengths.tail.scan(lengths.head)(_ + _)
   }
 }