Cache preferred locations of checkpointed RDD.

Author: viirya

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

@@ -783,6 +783,13 @@ package object config {
       .booleanConf
       .createWithDefault(false)
 
+  private[spark] val CACHE_CHECKPOINT_PREFERRED_LOCS =
+    ConfigBuilder("spark.rdd.checkpoint.cachePreferredLocs")
+      .internal()
+      .doc("Whether to cache preferred locations of checkpointed RDD.")
+      .booleanConf
+      .createWithDefault(false)
+
   private[spark] val SHUFFLE_ACCURATE_BLOCK_THRESHOLD =
     ConfigBuilder("spark.shuffle.accurateBlockThreshold")
       .doc("Threshold in bytes above which the size of shuffle blocks in " +

core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala

@@ -20,6 +20,7 @@ package org.apache.spark.rdd
 import java.io.{FileNotFoundException, IOException}
 import java.util.concurrent.TimeUnit
 
+import scala.collection.mutable
 import scala.reflect.ClassTag
 import scala.util.control.NonFatal
 
@@ -28,7 +29,7 @@ import org.apache.hadoop.fs.Path
 import org.apache.spark._
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.internal.Logging
-import org.apache.spark.internal.config.{BUFFER_SIZE, CHECKPOINT_COMPRESS}
+import org.apache.spark.internal.config.{BUFFER_SIZE, CACHE_CHECKPOINT_PREFERRED_LOCS, CHECKPOINT_COMPRESS}
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.util.{SerializableConfiguration, Utils}
 
@@ -82,14 +83,28 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
     Array.tabulate(inputFiles.length)(i => new CheckpointRDDPartition(i))
   }
 
+  // Cache of preferred locations of checkpointed files.
+  private[spark] val cachedPreferredLocations: mutable.HashMap[Int, Seq[String]] =
+    mutable.HashMap.empty
+
   /**
    * Return the locations of the checkpoint file associated with the given partition.
    */
   protected override def getPreferredLocations(split: Partition): Seq[String] = {
-    val status = fs.getFileStatus(
-      new Path(checkpointPath, ReliableCheckpointRDD.checkpointFileName(split.index)))
-    val locations = fs.getFileBlockLocations(status, 0, status.getLen)
-    locations.headOption.toList.flatMap(_.getHosts).filter(_ != "localhost")
+    val cachePreferrredLoc = SparkEnv.get.conf.get(CACHE_CHECKPOINT_PREFERRED_LOCS)
+
+    if (cachePreferrredLoc && cachedPreferredLocations.contains(split.index)) {
+      cachedPreferredLocations(split.index)
+    } else {
+      val status = fs.getFileStatus(
+        new Path(checkpointPath, ReliableCheckpointRDD.checkpointFileName(split.index)))
+      val locations = fs.getFileBlockLocations(status, 0, status.getLen)
+      val preferredLoc = locations.headOption.toList.flatMap(_.getHosts).filter(_ != "localhost")
+      if (cachePreferrredLoc) {
+        cachedPreferredLocations.update(split.index, preferredLoc)
+      }
+      preferredLoc
+    }
   }
 
   /**

core/src/test/scala/org/apache/spark/CheckpointSuite.scala

@@ -584,7 +584,7 @@ object CheckpointSuite {
   }
 }
 
-class CheckpointCompressionSuite extends SparkFunSuite with LocalSparkContext {
+class CheckpointStorageSuite extends SparkFunSuite with LocalSparkContext {
 
   test("checkpoint compression") {
     withTempDir { checkpointDir =>
@@ -618,4 +618,26 @@ class CheckpointCompressionSuite extends SparkFunSuite with LocalSparkContext {
       assert(rdd.collect().toSeq === (1 to 20))
     }
   }
+
+  test("cache checkpoint preferred location") {
+    withTempDir { checkpointDir =>
+      val conf = new SparkConf()
+        .set("spark.rdd.checkpoint.cachePreferredLocs", "true")
+        .set(UI_ENABLED.key, "false")
+      sc = new SparkContext("local", "test", conf)
+      sc.setCheckpointDir(checkpointDir.toString)
+      val rdd = sc.makeRDD(1 to 20, numSlices = 1)
+      rdd.checkpoint()
+      assert(rdd.collect().toSeq === (1 to 20))
+
+      // Verify that RDD is checkpointed
+      assert(rdd.firstParent.isInstanceOf[ReliableCheckpointRDD[_]])
+      val checkpointedRDD = rdd.firstParent.asInstanceOf[ReliableCheckpointRDD[_]]
+      assert(!checkpointedRDD.cachedPreferredLocations.isDefinedAt(0))
+
+      val preferredLoc = checkpointedRDD.preferredLocations(checkpointedRDD.partitions(0))
+      assert(checkpointedRDD.cachedPreferredLocations.isDefinedAt(0))
+      assert(preferredLoc == checkpointedRDD.cachedPreferredLocations(0))
+    }
+  }
 }