[SPARK-23253][Core][Shuffle]Only write shuffle temporary index file when there is not an existing one

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

@@ -141,19 +141,6 @@ private[spark] class IndexShuffleBlockResolver(
     val indexFile = getIndexFile(shuffleId, mapId)
     val indexTmp = Utils.tempFileWith(indexFile)
     try {
-      val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexTmp)))
-      Utils.tryWithSafeFinally {
-        // We take in lengths of each block, need to convert it to offsets.
-        var offset = 0L
-        out.writeLong(offset)
-        for (length <- lengths) {
-          offset += length
-          out.writeLong(offset)
-        }
-      } {
-        out.close()
-      }
-
       val dataFile = getDataFile(shuffleId, mapId)
       // There is only one IndexShuffleBlockResolver per executor, this synchronization make sure
       // the following check and rename are atomic.
@@ -166,10 +153,22 @@ private[spark] class IndexShuffleBlockResolver(
           if (dataTmp != null && dataTmp.exists()) {
             dataTmp.delete()
           }
-          indexTmp.delete()
         } else {
           // This is the first successful attempt in writing the map outputs for this task,
           // so override any existing index and data files with the ones we wrote.
+          val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexTmp)))
+          Utils.tryWithSafeFinally {
+            // We take in lengths of each block, need to convert it to offsets.
+            var offset = 0L
+            out.writeLong(offset)
+            for (length <- lengths) {
+              offset += length
+              out.writeLong(offset)
+            }
+          } {
+            out.close()
+          }
+
           if (indexFile.exists()) {
             indexFile.delete()
           }

core/src/test/scala/org/apache/spark/shuffle/sort/IndexShuffleBlockResolverSuite.scala

@@ -17,7 +17,7 @@
 
 package org.apache.spark.shuffle.sort
 
-import java.io.{File, FileInputStream, FileOutputStream}
+import java.io._
 
 import org.mockito.{Mock, MockitoAnnotations}
 import org.mockito.Answers.RETURNS_SMART_NULLS
@@ -123,7 +123,7 @@ class IndexShuffleBlockResolverSuite extends SparkFunSuite with BeforeAndAfterEa
     assert(dataFile.length() === 35)
     assert(!dataTmp2.exists())
 
-    // The dataFile should be the previous one
+    // The dataFile should be the new one, since we deleted the dataFile from the first attempt
     val firstByte2 = new Array[Byte](1)
     val in2 = new FileInputStream(dataFile)
     Utils.tryWithSafeFinally {
@@ -133,4 +133,65 @@ class IndexShuffleBlockResolverSuite extends SparkFunSuite with BeforeAndAfterEa
     }
     assert(firstByte2(0) === 2)
   }
+
+  test("SPARK-23253: index files should be created properly") {
+    val shuffleId = 1
+    val mapId = 2
+    val idxName = s"shuffle_${shuffleId}_${mapId}_0.index"
+    val resolver = new IndexShuffleBlockResolver(conf, blockManager)
+
+    val lengths = (1 to 2).map(_ => 8L).toArray
+    val dataTmp = File.createTempFile("shuffle", null, tempDir)
+    val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(dataTmp)))
+    Utils.tryWithSafeFinally {
+      lengths.foreach(out.writeLong)
+    } {
+      out.close()
+    }
+    resolver.writeIndexFileAndCommit(shuffleId, mapId, lengths, dataTmp)
+
+    val indexFile = new File(tempDir.getAbsolutePath, idxName)
+    val dataFile = resolver.getDataFile(shuffleId, mapId)
+    val idxFile1Len = indexFile.length()
+
+    assert(indexFile.exists())
+    assert(idxFile1Len === (lengths.length + 1) * 8)
+    assert(dataFile.exists())
+    assert(dataFile.length() === 8 * lengths.length)
+    assert(!dataTmp.exists())
+
+    // delete dataFile, index file will be replaced.
+    dataFile.delete()
+    assert(!dataFile.exists())
+
+    val lengths2 = (1 to 4).map(_ => 8L).toArray
+    val dataTmp2 = File.createTempFile("shuffle", null, tempDir)
+    val out2 = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(dataTmp2)))
+    Utils.tryWithSafeFinally {
+      lengths2.foreach(out2.writeLong)
+    } {
+      out2.close()
+    }
+    resolver.writeIndexFileAndCommit(shuffleId, mapId, lengths2, dataTmp2)
+    val idxFile2Len = indexFile.length()
+
+    assert(indexFile.exists())
+    assert(indexFile.length() === (lengths2.length + 1) * 8)
+    assert(idxFile1Len !== idxFile2Len, "index file should be updated.")
+
+    // all files are present and will be reused
+    val lengths3 = (1 to 6).map(_ => 8L).toArray
+    val dataTmp3 = File.createTempFile("shuffle", null, tempDir)
+    val out3 = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(dataTmp3)))
+    Utils.tryWithSafeFinally {
+      lengths3.foreach(out3.writeLong)
+    } {
+      out3.close()
+    }
+    resolver.writeIndexFileAndCommit(shuffleId, mapId, lengths2, dataTmp2)
+    val idxFile3Len = indexFile.length()
+
+    assert(indexFile.exists())
+    assert(idxFile2Len === idxFile3Len, "index file should not change.")
+  }
 }