[SPARK-25299] Move shuffle writers back to being given specific partition ids (apache-spark-on-k8s#540)

We originally made the shuffle map output writer API behave like an iterator in fetching the "next" partition writer. However, the shuffle writer implementations tend to skip opening empty partitions. If we used an iterator-like API though we would be tied down to opening a partition writer for every single partition, even if some of them are empty. Here, we go back to using specific partition identifiers to give us more freedom to avoid needing to create writers for empty partitions.

Author: mccheah

core/src/main/java/org/apache/spark/api/shuffle/ShuffleMapOutputWriter.java

@@ -30,7 +30,7 @@
  */
 @Experimental
 public interface ShuffleMapOutputWriter {
-  ShufflePartitionWriter getNextPartitionWriter() throws IOException;
+  ShufflePartitionWriter getPartitionWriter(int partitionId) throws IOException;
 
   Optional<MapShuffleLocations> commitAllPartitions() throws IOException;
 

core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java

@@ -216,7 +216,7 @@ private long[] writePartitionedData(ShuffleMapOutputWriter mapOutputWriter) thro
         boolean copyThrewException = true;
         ShufflePartitionWriter writer = null;
         try {
-          writer = mapOutputWriter.getNextPartitionWriter();
+          writer = mapOutputWriter.getPartitionWriter(i);
           if (!file.exists()) {
             copyThrewException = false;
           } else {

core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java

@@ -291,18 +291,6 @@ private long[] mergeSpills(SpillInfo[] spills,
     long[] partitionLengths = new long[numPartitions];
     try {
       if (spills.length == 0) {
-        // The contract we are working under states that we will open a partition writer for
-        // each partition, regardless of number of spills
-        for (int i = 0; i < numPartitions; i++) {
-          ShufflePartitionWriter writer = null;
-          try {
-            writer = mapWriter.getNextPartitionWriter();
-          } finally {
-            if (writer != null) {
-              writer.close();
-            }
-          }
-        }
         return partitionLengths;
       } else {
         // There are multiple spills to merge, so none of these spill files' lengths were counted
@@ -378,7 +366,7 @@ private long[] mergeSpillsWithFileStream(
         boolean copyThrewExecption = true;
         ShufflePartitionWriter writer = null;
         try {
-          writer = mapWriter.getNextPartitionWriter();
+          writer = mapWriter.getPartitionWriter(partition);
           OutputStream partitionOutput = null;
           try {
             // Shield the underlying output stream from close() calls, so that we can close the
@@ -457,7 +445,7 @@ private long[] mergeSpillsWithTransferTo(
         boolean copyThrewExecption = true;
         ShufflePartitionWriter writer = null;
         try {
-          writer = mapWriter.getNextPartitionWriter();
+          writer = mapWriter.getPartitionWriter(partition);
           WritableByteChannel channel = writer.toChannel();
           for (int i = 0; i < spills.length; i++) {
             long partitionLengthInSpill = 0L;

core/src/main/java/org/apache/spark/shuffle/sort/io/DefaultShuffleMapOutputWriter.java

@@ -51,7 +51,7 @@ public class DefaultShuffleMapOutputWriter implements ShuffleMapOutputWriter {
   private final IndexShuffleBlockResolver blockResolver;
   private final long[] partitionLengths;
   private final int bufferSize;
-  private int currPartitionId = 0;
+  private int lastPartitionId = -1;
   private long currChannelPosition;
   private final BlockManagerId shuffleServerId;
 
@@ -84,7 +84,11 @@ public DefaultShuffleMapOutputWriter(
   }
 
   @Override
-  public ShufflePartitionWriter getNextPartitionWriter() throws IOException {
+  public ShufflePartitionWriter getPartitionWriter(int partitionId) throws IOException {
+    if (partitionId <= lastPartitionId) {
+      throw new IllegalArgumentException("Partitions should be requested in increasing order.");
+    }
+    lastPartitionId = partitionId;
     if (outputTempFile == null) {
       outputTempFile = Utils.tempFileWith(outputFile);
     }
@@ -93,7 +97,7 @@ public ShufflePartitionWriter getNextPartitionWriter() throws IOException {
     } else {
       currChannelPosition = 0L;
     }
-    return new DefaultShufflePartitionWriter(currPartitionId++);
+    return new DefaultShufflePartitionWriter(partitionId);
   }
 
   @Override

core/src/main/scala/org/apache/spark/util/collection/ExternalSorter.scala

@@ -717,17 +717,6 @@ private[spark] class ExternalSorter[K, V, C](
     lengths
   }
 
-  private def writeEmptyPartition(mapOutputWriter: ShuffleMapOutputWriter): Unit = {
-    var partitionWriter: ShufflePartitionWriter = null
-    try {
-      partitionWriter = mapOutputWriter.getNextPartitionWriter
-    } finally {
-      if (partitionWriter != null) {
-        partitionWriter.close()
-      }
-    }
-  }
-
   /**
    * Write all the data added into this ExternalSorter into a map output writer that pushes bytes
    * to some arbitrary backing store. This is called by the SortShuffleWriter.
@@ -738,26 +727,16 @@ private[spark] class ExternalSorter[K, V, C](
       shuffleId: Int, mapId: Int, mapOutputWriter: ShuffleMapOutputWriter): Array[Long] = {
     // Track location of each range in the map output
     val lengths = new Array[Long](numPartitions)
-    var nextPartitionId = 0
     if (spills.isEmpty) {
       // Case where we only have in-memory data
       val collection = if (aggregator.isDefined) map else buffer
       val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
       while (it.hasNext()) {
         val partitionId = it.nextPartition()
-        // The contract for the plugin is that we will ask for a writer for every partition
-        // even if it's empty. However, the external sorter will return non-contiguous
-        // partition ids. So this loop "backfills" the empty partitions that form the gaps.
-
-        // The algorithm as a whole is correct because the partition ids are returned by the
-        // iterator in ascending order.
-        for (emptyPartition <- nextPartitionId until partitionId) {
-          writeEmptyPartition(mapOutputWriter)
-        }
         var partitionWriter: ShufflePartitionWriter = null
         var partitionPairsWriter: ShufflePartitionPairsWriter = null
         try {
-          partitionWriter = mapOutputWriter.getNextPartitionWriter
+          partitionWriter = mapOutputWriter.getPartitionWriter(partitionId)
           val blockId = ShuffleBlockId(shuffleId, mapId, partitionId)
           partitionPairsWriter = new ShufflePartitionPairsWriter(
             partitionWriter,
@@ -779,7 +758,6 @@ private[spark] class ExternalSorter[K, V, C](
         if (partitionWriter != null) {
           lengths(partitionId) = partitionWriter.getNumBytesWritten
         }
-        nextPartitionId = partitionId + 1
       }
     } else {
       // We must perform merge-sort; get an iterator by partition and write everything directly.
@@ -790,14 +768,11 @@ private[spark] class ExternalSorter[K, V, C](
 
         // The algorithm as a whole is correct because the partition ids are returned by the
         // iterator in ascending order.
-        for (emptyPartition <- nextPartitionId until id) {
-          writeEmptyPartition(mapOutputWriter)
-        }
         val blockId = ShuffleBlockId(shuffleId, mapId, id)
         var partitionWriter: ShufflePartitionWriter = null
         var partitionPairsWriter: ShufflePartitionPairsWriter = null
         try {
-          partitionWriter = mapOutputWriter.getNextPartitionWriter
+          partitionWriter = mapOutputWriter.getPartitionWriter(id)
           partitionPairsWriter = new ShufflePartitionPairsWriter(
             partitionWriter,
             serializerManager,
@@ -817,16 +792,9 @@ private[spark] class ExternalSorter[K, V, C](
         if (partitionWriter != null) {
           lengths(id) = partitionWriter.getNumBytesWritten
         }
-        nextPartitionId = id + 1
       }
     }
 
-    // The iterator may have stopped short of opening a writer for every partition. So fill in the
-    // remaining empty partitions.
-    for (emptyPartition <- nextPartitionId until numPartitions) {
-      writeEmptyPartition(mapOutputWriter)
-    }
-
     context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
     context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
     context.taskMetrics().incPeakExecutionMemory(peakMemoryUsedBytes)

core/src/test/scala/org/apache/spark/shuffle/sort/io/DefaultShuffleMapOutputWriterSuite.scala

@@ -140,7 +140,7 @@ class DefaultShuffleMapOutputWriterSuite extends SparkFunSuite with BeforeAndAft
 
   test("writing to an outputstream") {
     (0 until NUM_PARTITIONS).foreach{ p =>
-      val writer = mapOutputWriter.getNextPartitionWriter
+      val writer = mapOutputWriter.getPartitionWriter(p)
       val stream = writer.toStream()
       data(p).foreach { i => stream.write(i)}
       stream.close()
@@ -159,7 +159,7 @@ class DefaultShuffleMapOutputWriterSuite extends SparkFunSuite with BeforeAndAft
 
   test("writing to a channel") {
     (0 until NUM_PARTITIONS).foreach{ p =>
-      val writer = mapOutputWriter.getNextPartitionWriter
+      val writer = mapOutputWriter.getPartitionWriter(p)
       val channel = writer.toChannel()
       val byteBuffer = ByteBuffer.allocate(D_LEN * 4)
       val intBuffer = byteBuffer.asIntBuffer()
@@ -179,7 +179,7 @@ class DefaultShuffleMapOutputWriterSuite extends SparkFunSuite with BeforeAndAft
 
   test("copyStreams with an outputstream") {
     (0 until NUM_PARTITIONS).foreach{ p =>
-      val writer = mapOutputWriter.getNextPartitionWriter
+      val writer = mapOutputWriter.getPartitionWriter(p)
       val stream = writer.toStream()
       val byteBuffer = ByteBuffer.allocate(D_LEN * 4)
       val intBuffer = byteBuffer.asIntBuffer()
@@ -200,7 +200,7 @@ class DefaultShuffleMapOutputWriterSuite extends SparkFunSuite with BeforeAndAft
 
   test("copyStreamsWithNIO with a channel") {
     (0 until NUM_PARTITIONS).foreach{ p =>
-      val writer = mapOutputWriter.getNextPartitionWriter
+      val writer = mapOutputWriter.getPartitionWriter(p)
       val channel = writer.toChannel()
       val byteBuffer = ByteBuffer.allocate(D_LEN * 4)
       val intBuffer = byteBuffer.asIntBuffer()