[SPARK-11051][Core] Do not allow local checkpointing after the RDD is materialized and checkpointed

core/src/main/scala/org/apache/spark/api/java/JavaRDDLike.scala

@@ -547,6 +547,11 @@ trait JavaRDDLike[T, This <: JavaRDDLike[T, This]] extends Serializable {
    */
   def isCheckpointed: Boolean = rdd.isCheckpointed
 
+  /**
+   * Return whether this RDD has been checkpointed and materialized or not
+   */
+  private[spark] def isCheckpointedAndMaterialized: Boolean = rdd.isCheckpointedAndMaterialized
+
   /**
    * Gets the name of the file to which this RDD was checkpointed
    */

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

@@ -294,7 +294,11 @@ abstract class RDD[T: ClassTag](
    */
   private[spark] def computeOrReadCheckpoint(split: Partition, context: TaskContext): Iterator[T] =
   {
-    if (isCheckpointed) firstParent[T].iterator(split, context) else compute(split, context)
+    if (isCheckpointedAndMaterialized) {
+      firstParent[T].iterator(split, context)
+    } else {
+      compute(split, context)
+    }
   }
 
   /**
@@ -1520,12 +1524,22 @@ abstract class RDD[T: ClassTag](
       persist(LocalRDDCheckpointData.transformStorageLevel(storageLevel), allowOverride = true)
     }
 
-    checkpointData match {
-      case Some(reliable: ReliableRDDCheckpointData[_]) => logWarning(
-        "RDD was already marked for reliable checkpointing: overriding with local checkpoint.")
-      case _ =>
+    // If this RDD is already checkpointed and materialized, its lineage is already truncated.
+    // We must not override our `checkpointData` in this case because it is needed to recover
+    // the checkpointed data. If it is overridden, next time materializing on this RDD will
+    // cause error.
+    if (isCheckpointedAndMaterialized) {
+      logWarning("Not marking RDD for local checkpoint because it was already " +
+        "checkpointed and materialized")
+    } else {
+      // Lineage is not truncated yet, so just override any existing checkpoint data with ours
+      checkpointData match {
+        case Some(_: ReliableRDDCheckpointData[_]) => logWarning(
+          "RDD was already marked for reliable checkpointing: overriding with local checkpoint.")
+        case _ =>
+      }
+      checkpointData = Some(new LocalRDDCheckpointData(this))
     }
-    checkpointData = Some(new LocalRDDCheckpointData(this))
     this
   }
 
@@ -1534,6 +1548,12 @@ abstract class RDD[T: ClassTag](
    */
   def isCheckpointed: Boolean = checkpointData.exists(_.isCheckpointed)
 
+  /**
+   * Return whether this RDD is marked for checkpointing and materialized,
+   * either reliably or locally.
+   */
+  private[spark] def isCheckpointedAndMaterialized: Boolean = isCheckpointed
+
   /**
    * Return whether this RDD is marked for local checkpointing.
    * Exposed for testing.

core/src/test/java/org/apache/spark/JavaAPISuite.java

@@ -1420,9 +1420,13 @@ public void checkpointAndComputation() {
     JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4, 5));
     sc.setCheckpointDir(tempDir.getAbsolutePath());
     Assert.assertFalse(rdd.isCheckpointed());
+    Assert.assertFalse(rdd.isCheckpointedAndMaterialized());
     rdd.checkpoint();
+    Assert.assertFalse(rdd.isCheckpointed());
+    Assert.assertFalse(rdd.isCheckpointedAndMaterialized());
     rdd.count(); // Forces the DAG to cause a checkpoint
     Assert.assertTrue(rdd.isCheckpointed());
+    Assert.assertTrue(rdd.isCheckpointedAndMaterialized());
     Assert.assertEquals(Arrays.asList(1, 2, 3, 4, 5), rdd.collect());
   }
 
@@ -1431,9 +1435,13 @@ public void checkpointAndRestore() {
     JavaRDD<Integer> rdd = sc.parallelize(Arrays.asList(1, 2, 3, 4, 5));
     sc.setCheckpointDir(tempDir.getAbsolutePath());
     Assert.assertFalse(rdd.isCheckpointed());
+    Assert.assertFalse(rdd.isCheckpointedAndMaterialized());
     rdd.checkpoint();
+    Assert.assertFalse(rdd.isCheckpointed());
+    Assert.assertFalse(rdd.isCheckpointedAndMaterialized());
     rdd.count(); // Forces the DAG to cause a checkpoint
     Assert.assertTrue(rdd.isCheckpointed());
+    Assert.assertTrue(rdd.isCheckpointedAndMaterialized());
 
     Assert.assertTrue(rdd.getCheckpointFile().isPresent());
     JavaRDD<Integer> recovered = sc.checkpointFile(rdd.getCheckpointFile().get());

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

@@ -241,9 +241,13 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     val rdd = new BlockRDD[Int](sc, Array[BlockId]())
     assert(rdd.partitions.size === 0)
     assert(rdd.isCheckpointed === false)
+    assert(rdd.isCheckpointedAndMaterialized === false)
     checkpoint(rdd, reliableCheckpoint)
+    assert(rdd.isCheckpointed === false)
+    assert(rdd.isCheckpointedAndMaterialized === false)
     assert(rdd.count() === 0)
     assert(rdd.isCheckpointed === true)
+    assert(rdd.isCheckpointedAndMaterialized === true)
     assert(rdd.partitions.size === 0)
   }