[SPARK-32003] Update fileLostEpoch.

I inadvertently left out a line when transferring code. The fileLostEpoch
needs to be updated with an entry for the failed executor with lost output.

Adopted suggestions from wuyi and attilapiros.

Author: wypoon

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

@@ -1944,16 +1944,17 @@ private[spark] class DAGScheduler(
     logDebug(s"Removing executor $execId, fileLost: $fileLost, currentEpoch: $currentEpoch")
     if (!failedEpoch.contains(execId) || failedEpoch(execId) < currentEpoch) {
       failedEpoch(execId) = currentEpoch
-      logInfo("Executor lost: %s (epoch %d)".format(execId, currentEpoch))
+      logInfo(s"Executor lost: $execId (epoch $currentEpoch)")
       blockManagerMaster.removeExecutor(execId)
     }
     if (fileLost && (!fileLostEpoch.contains(execId) || fileLostEpoch(execId) < currentEpoch)) {
+      fileLostEpoch(execId) = currentEpoch
       hostToUnregisterOutputs match {
         case Some(host) =>
-          logInfo("Shuffle files lost for host: %s (epoch %d)".format(host, currentEpoch))
+          logInfo(s"Shuffle files lost for host: $host (epoch $currentEpoch)")
           mapOutputTracker.removeOutputsOnHost(host)
         case None =>
-          logInfo("Shuffle files lost for executor: %s (epoch %d)".format(execId, currentEpoch))
+          logInfo(s"Shuffle files lost for executor: $execId (epoch $currentEpoch)")
           mapOutputTracker.removeOutputsOnExecutor(execId)
       }
       clearCacheLocs()
@@ -1986,9 +1987,7 @@ private[spark] class DAGScheduler(
       logInfo("Host added was in lost list earlier: " + host)
       failedEpoch -= execId
     }
-    if (fileLostEpoch.contains(execId)) {
-      fileLostEpoch -= execId
-    }
+    fileLostEpoch -= execId
   }
 
   private[scheduler] def handleStageCancellation(stageId: Int, reason: Option[String]): Unit = {

core/src/test/scala/org/apache/spark/scheduler/DAGSchedulerSuite.scala

@@ -555,19 +555,15 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi
     submit(reduceRdd, Array(0, 1, 2))
     // Map stage completes successfully,
     // two tasks are run on an executor on hostA and one on an executor on hostB
-    complete(taskSets(0), Seq(
-      (Success, makeMapStatus("hostA", 3)),
-      (Success, makeMapStatus("hostA", 3)),
-      (Success, makeMapStatus("hostB", 3))))
+    completeShuffleMapStageSuccessfully(0, 0, 3, Seq("hostA", "hostA", "hostB"))
     // Now the executor on hostA is lost
     runEvent(ExecutorLost("hostA-exec", ExecutorExited(-100, false, "Container marked as failed")))
 
     // The MapOutputTracker has all the shuffle files
     val initialMapStatuses = mapOutputTracker.shuffleStatuses(shuffleId).mapStatuses
-    assert(initialMapStatuses.count(_ != null) == 3)
-    assert(initialMapStatuses(0).location.executorId === "hostA-exec")
-    assert(initialMapStatuses(1).location.executorId === "hostA-exec")
-    assert(initialMapStatuses(2).location.executorId === "hostB-exec")
+    assert(initialMapStatuses.count(_ != null) === 3)
+    assert(initialMapStatuses.count(s => s != null && s.location.executorId == "hostA-exec") === 2)
+    assert(initialMapStatuses.count(s => s != null && s.location.executorId == "hostB-exec") === 1)
 
     // Now a fetch failure from the lost executor occurs
     complete(taskSets(1), Seq(
@@ -576,8 +572,9 @@ class DAGSchedulerSuite extends SparkFunSuite with LocalSparkContext with TimeLi
 
     // Shuffle files for hostA-exec should be lost
     val mapStatuses = mapOutputTracker.shuffleStatuses(shuffleId).mapStatuses
-    assert(mapStatuses.count(_ != null) == 1)
-    assert(mapStatuses(2).location.executorId === "hostB-exec")
+    assert(mapStatuses.count(_ != null) === 1)
+    assert(initialMapStatuses.count(s => s != null && s.location.executorId == "hostA-exec") === 0)
+    assert(initialMapStatuses.count(s => s != null && s.location.executorId == "hostB-exec") === 1)
   }
 
   test("zero split job") {