containsBlock/removeDataByMap

Author: cxzl25

core/src/main/scala/org/apache/spark/shuffle/IndexShuffleBlockResolver.scala

@@ -62,7 +62,11 @@ private[spark] class IndexShuffleBlockResolver(
   private val remoteShuffleMaxDisk: Option[Long] =
     conf.get(config.STORAGE_DECOMMISSION_SHUFFLE_MAX_DISK_SIZE)
 
-  def getDataFile(shuffleId: Int, mapId: Long): File = getDataFile(shuffleId, mapId, None)
+  def getDataFile(shuffleId: Int, mapId: Long): File =
+    getDataFile(shuffleId, mapId, None, true)
+
+  def getDataFile(shuffleId: Int, mapId: Long, needCreate: Boolean): File =
+    getDataFile(shuffleId, mapId, None, needCreate)
 
   /**
    * Get the shuffle files that are stored locally. Used for block migrations.
@@ -95,12 +99,16 @@ private[spark] class IndexShuffleBlockResolver(
    * When the dirs parameter is None then use the disk manager's local directories. Otherwise,
    * read from the specified directories.
    */
-   def getDataFile(shuffleId: Int, mapId: Long, dirs: Option[Array[String]]): File = {
+  def getDataFile(
+      shuffleId: Int,
+      mapId: Long,
+      dirs: Option[Array[String]],
+      needCreate: Boolean): File = {
     val blockId = ShuffleDataBlockId(shuffleId, mapId, NOOP_REDUCE_ID)
     dirs
       .map(d =>
         new File(ExecutorDiskUtils.getFilePath(d, blockManager.subDirsPerLocalDir, blockId.name)))
-      .getOrElse(blockManager.diskBlockManager.getFile(blockId))
+      .getOrElse(blockManager.diskBlockManager.getFile(blockId, needCreate))
   }
 
   /**
@@ -112,12 +120,13 @@ private[spark] class IndexShuffleBlockResolver(
   def getIndexFile(
       shuffleId: Int,
       mapId: Long,
-      dirs: Option[Array[String]] = None): File = {
+      dirs: Option[Array[String]] = None,
+      needCreate: Boolean = true): File = {
     val blockId = ShuffleIndexBlockId(shuffleId, mapId, NOOP_REDUCE_ID)
     dirs
       .map(d =>
         new File(ExecutorDiskUtils.getFilePath(d, blockManager.subDirsPerLocalDir, blockId.name)))
-      .getOrElse(blockManager.diskBlockManager.getFile(blockId))
+      .getOrElse(blockManager.diskBlockManager.getFile(blockId, needCreate))
   }
 
   private def getMergedBlockDataFile(
@@ -154,17 +163,18 @@ private[spark] class IndexShuffleBlockResolver(
    * Remove data file and index file that contain the output data from one map.
    */
   def removeDataByMap(shuffleId: Int, mapId: Long): Unit = {
-    var file = getDataFile(shuffleId, mapId)
+    var file = getDataFile(shuffleId, mapId, needCreate = false)
     if (file.exists() && !file.delete()) {
       logWarning(s"Error deleting data ${file.getPath()}")
     }
 
-    file = getIndexFile(shuffleId, mapId)
+    file = getIndexFile(shuffleId, mapId, needCreate = false)
     if (file.exists() && !file.delete()) {
       logWarning(s"Error deleting index ${file.getPath()}")
     }
 
-    file = getChecksumFile(shuffleId, mapId, conf.get(config.SHUFFLE_CHECKSUM_ALGORITHM))
+    file = getChecksumFile(shuffleId, mapId, conf.get(config.SHUFFLE_CHECKSUM_ALGORITHM),
+      needCreate = false)
     if (file.exists() && !file.delete()) {
       logWarning(s"Error deleting checksum ${file.getPath()}")
     }
@@ -549,13 +559,14 @@ private[spark] class IndexShuffleBlockResolver(
       shuffleId: Int,
       mapId: Long,
       algorithm: String,
-      dirs: Option[Array[String]] = None): File = {
+      dirs: Option[Array[String]] = None,
+      needCreate: Boolean = true): File = {
     val blockId = ShuffleChecksumBlockId(shuffleId, mapId, NOOP_REDUCE_ID)
     val fileName = ShuffleChecksumHelper.getChecksumFileName(blockId.name, algorithm)
     dirs
       .map(d =>
         new File(ExecutorDiskUtils.getFilePath(d, blockManager.subDirsPerLocalDir, fileName)))
-      .getOrElse(blockManager.diskBlockManager.getFile(fileName))
+      .getOrElse(blockManager.diskBlockManager.getFile(fileName, needCreate))
   }
 
   override def getBlockData(
@@ -594,7 +605,7 @@ private[spark] class IndexShuffleBlockResolver(
       }
       new FileSegmentManagedBuffer(
         transportConf,
-        getDataFile(shuffleId, mapId, dirs),
+        getDataFile(shuffleId, mapId, dirs, true),
         startOffset,
         endOffset - startOffset)
     } finally {

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -302,7 +302,8 @@ private[spark] class BlockManager(
     val shuffleBlock = blockId.asInstanceOf[ShuffleBlockId]
     val resolver = shuffleManager.shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver]
     val checksumFile =
-      resolver.getChecksumFile(shuffleBlock.shuffleId, shuffleBlock.mapId, algorithm)
+      resolver.getChecksumFile(shuffleBlock.shuffleId, shuffleBlock.mapId, algorithm,
+        needCreate = false)
     val reduceId = shuffleBlock.reduceId
     ShuffleChecksumHelper.diagnoseCorruption(
       algorithm, checksumFile, reduceId, resolver.getBlockData(shuffleBlock), checksumByReader)

core/src/main/scala/org/apache/spark/storage/DiskBlockManager.scala

@@ -91,10 +91,10 @@ private[spark] class DiskBlockManager(
   // This method should be kept in sync with
   // org.apache.spark.network.shuffle.ExecutorDiskUtils#getFilePath().
   def getFile(filename: String): File = {
-    getFile(filename, true)
+    getFile(filename, needCreate = true)
   }
 
-  def getFile(filename: String, needCreate: Boolean = true): File = {
+  def getFile(filename: String, needCreate: Boolean): File = {
     // Figure out which local directory it hashes to, and which subdirectory in that
     val hash = Utils.nonNegativeHash(filename)
     val dirId = hash % localDirs.length
@@ -131,7 +131,10 @@ private[spark] class DiskBlockManager(
     new File(subDir, filename)
   }
 
-  def getFile(blockId: BlockId): File = getFile(blockId.name)
+  def getFile(blockId: BlockId): File = getFile(blockId.name, needCreate = true)
+
+  def getFile(blockId: BlockId, needCreate: Boolean): File =
+    getFile(blockId.name, needCreate)
 
   /**
    * This should be in sync with
@@ -162,7 +165,7 @@ private[spark] class DiskBlockManager(
 
   /** Check if disk block manager has a block. */
   def containsBlock(blockId: BlockId): Boolean = {
-    getFile(blockId.name).exists()
+    getFile(blockId.name, needCreate = false).exists()
   }
 
   /** List all the files currently stored on disk by the disk manager. */

core/src/main/scala/org/apache/spark/storage/DiskStore.scala

@@ -128,7 +128,7 @@ private[spark] class DiskStore(
 
   def remove(blockId: BlockId): Boolean = {
     blockSizes.remove(blockId)
-    val file = diskManager.getFile(blockId.name, false)
+    val file = diskManager.getFile(blockId.name, needCreate = false)
     if (file.exists()) {
       val ret = file.delete()
       if (!ret) {