Address more comments.

Author: mccheah

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

@@ -18,6 +18,7 @@
 package org.apache.spark.shuffle.api;
 
 import java.io.IOException;
+
 import org.apache.spark.annotation.Private;
 
 /**

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

@@ -18,11 +18,10 @@
 package org.apache.spark.shuffle.api;
 
 import java.io.IOException;
+import java.util.Optional;
 import java.io.OutputStream;
-import java.nio.channels.Channels;
 
 import org.apache.spark.annotation.Private;
-import org.apache.spark.shuffle.sort.io.DefaultWritableByteChannelWrapper;
 
 /**
  * :: Private ::
@@ -67,25 +66,24 @@ public interface ShufflePartitionWriter {
    * Implementations that intend on combining the bytes for all the partitions written by this
    * map task should reuse the same channel instance across all the partition writers provided
    * by the parent {@link ShuffleMapOutputWriter}. If one does so, ensure that
-   * {@link WritableByteChannelWrapper#close()} does not close the resource, since it
+   * {@link WritableByteChannelWrapper#close()} does not close the resource, since the channel
    * will be reused across partition writes. The underlying resources should be cleaned up in
    * {@link ShuffleMapOutputWriter#commitAllPartitions()} and
    * {@link ShuffleMapOutputWriter#abort(Throwable)}.
    * <p>
    * This method is primarily for advanced optimizations where bytes can be copied from the input
-   * spill files to the output channel without copying data into memory.
-   * <p>
-   * The default implementation should be sufficient for most situations. Only override this
-   * method if there is a very specific optimization that needs to be built.
+   * spill files to the output channel without copying data into memory. If such optimizations are
+   * not supported, the implementation should return {@link Optional#empty()}. By default, the
+   * implementation returns {@link Optional#empty()}.
    * <p>
    * Note that the returned {@link WritableByteChannelWrapper} itself is closed, but not the
    * underlying channel that is returned by {@link WritableByteChannelWrapper#channel()}. Ensure
    * that the underlying channel is cleaned up in {@link WritableByteChannelWrapper#close()},
    * {@link ShuffleMapOutputWriter#commitAllPartitions()}, or
    * {@link ShuffleMapOutputWriter#abort(Throwable)}.
    */
-  default WritableByteChannelWrapper openChannelWrapper() throws IOException {
-    return new DefaultWritableByteChannelWrapper(Channels.newChannel(openStream()));
+  default Optional<WritableByteChannelWrapper> openChannelWrapper() throws IOException {
+    return Optional.empty();
   }
 
   /**

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

@@ -22,6 +22,7 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.channels.FileChannel;
+import java.util.Optional;
 import javax.annotation.Nullable;
 
 import scala.None$;
@@ -205,45 +206,23 @@ private long[] writePartitionedData(ShuffleMapOutputWriter mapOutputWriter) thro
         final File file = partitionWriterSegments[i].file();
         ShufflePartitionWriter writer = mapOutputWriter.getPartitionWriter(i);
         if (file.exists()) {
-          boolean copyThrewException = true;
           if (transferToEnabled) {
-            FileInputStream in = new FileInputStream(file);
             // Using WritableByteChannelWrapper to make resource closing consistent between
             // this implementation and UnsafeShuffleWriter.
-            try {
-              WritableByteChannelWrapper outputChannel = writer.openChannelWrapper();
-              try (FileChannel inputChannel = in.getChannel()) {
-                Utils.copyFileStreamNIO(
-                    inputChannel, outputChannel.channel(), 0L, inputChannel.size());
-                copyThrewException = false;
-              } finally {
-                Closeables.close(outputChannel, copyThrewException);
-              }
-            } finally {
-              Closeables.close(in, copyThrewException);
+            Optional<WritableByteChannelWrapper> maybeOutputChannel = writer.openChannelWrapper();
+            if (maybeOutputChannel.isPresent()) {
+              writePartitionedDataWithChannel(file, maybeOutputChannel.get());
+            } else {
+              writePartitionedDataWithStream(file, writer);
             }
           } else {
-            FileInputStream in = new FileInputStream(file);
-            OutputStream outputStream;
-            try {
-              outputStream = writer.openStream();
-              try {
-                Utils.copyStream(in, outputStream, false, false);
-                copyThrewException = false;
-              } finally {
-                Closeables.close(outputStream, copyThrewException);
-              }
-            } finally {
-              Closeables.close(in, copyThrewException);
-            }
+            writePartitionedDataWithStream(file, writer);
           }
           if (!file.delete()) {
             logger.error("Unable to delete file for partition {}", i);
           }
         }
-        long numBytesWritten = writer.getNumBytesWritten();
-        lengths[i] = numBytesWritten;
-        writeMetrics.incBytesWritten(numBytesWritten);
+        lengths[i] = writer.getNumBytesWritten();
       }
     } finally {
       writeMetrics.incWriteTime(System.nanoTime() - writeStartTime);
@@ -252,6 +231,41 @@ private long[] writePartitionedData(ShuffleMapOutputWriter mapOutputWriter) thro
     return lengths;
   }
 
+  private void writePartitionedDataWithChannel(
+      File file, WritableByteChannelWrapper outputChannel) throws IOException {
+    boolean copyThrewException = true;
+    try {
+      FileInputStream in = new FileInputStream(file);
+      try (FileChannel inputChannel = in.getChannel()) {
+        Utils.copyFileStreamNIO(
+            inputChannel, outputChannel.channel(), 0L, inputChannel.size());
+        copyThrewException = false;
+      } finally {
+        Closeables.close(in, copyThrewException);
+      }
+    } finally {
+      Closeables.close(outputChannel, copyThrewException);
+    }
+  }
+
+  private void writePartitionedDataWithStream(File file, ShufflePartitionWriter writer)
+      throws IOException {
+    boolean copyThrewException = true;
+    FileInputStream in = new FileInputStream(file);
+    OutputStream outputStream;
+    try {
+      outputStream = writer.openStream();
+      try {
+        Utils.copyStream(in, outputStream, false, false);
+        copyThrewException = false;
+      } finally {
+        Closeables.close(outputStream, copyThrewException);
+      }
+    } finally {
+      Closeables.close(in, copyThrewException);
+    }
+  }
+
   @Override
   public Option<MapStatus> stop(boolean success) {
     if (stopping) {

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

@@ -25,6 +25,7 @@
 import java.nio.channels.FileChannel;
 import java.nio.channels.WritableByteChannel;
 
+import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -164,7 +165,7 @@ public OutputStream openStream() throws IOException {
     }
 
     @Override
-    public WritableByteChannelWrapper openChannelWrapper() throws IOException {
+    public Optional<WritableByteChannelWrapper> openChannelWrapper() throws IOException {
       if (partChannel == null) {
         if (partStream != null) {
           throw new IllegalStateException("Requested an output stream for a previous write but" +
@@ -174,7 +175,7 @@ public WritableByteChannelWrapper openChannelWrapper() throws IOException {
         initChannel();
         partChannel = new PartitionWriterChannel(partitionId);
       }
-      return partChannel;
+      return Optional.of(partChannel);
     }
 
     @Override

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

@@ -17,7 +17,7 @@
 
 package org.apache.spark.shuffle.sort.io
 
-import java.io.{File, FileInputStream, FileOutputStream}
+import java.io.{File, FileInputStream}
 import java.nio.channels.FileChannel
 import java.nio.file.Files
 import java.util.Arrays
@@ -30,7 +30,6 @@ import org.mockito.MockitoAnnotations
 import org.scalatest.BeforeAndAfterEach
 
 import org.apache.spark.{SparkConf, SparkFunSuite}
-import org.apache.spark.executor.ShuffleWriteMetrics
 import org.apache.spark.shuffle.IndexShuffleBlockResolver
 import org.apache.spark.util.Utils
 
@@ -39,9 +38,6 @@ class LocalDiskShuffleMapOutputWriterSuite extends SparkFunSuite with BeforeAndA
   @Mock(answer = RETURNS_SMART_NULLS)
   private var blockResolver: IndexShuffleBlockResolver = _
 
-  @Mock(answer = RETURNS_SMART_NULLS)
-  private var shuffleWriteMetrics: ShuffleWriteMetrics = _
-
   private val NUM_PARTITIONS = 4
   private val data: Array[Array[Byte]] = (0 until NUM_PARTITIONS).map { p =>
     if (p == 3) {
@@ -93,7 +89,6 @@ class LocalDiskShuffleMapOutputWriterSuite extends SparkFunSuite with BeforeAndA
       0,
       0,
       NUM_PARTITIONS,
-      shuffleWriteMetrics,
       blockResolver,
       conf)
   }
@@ -116,13 +111,11 @@ class LocalDiskShuffleMapOutputWriterSuite extends SparkFunSuite with BeforeAndA
     (0 until NUM_PARTITIONS).foreach { p =>
       val writer = mapOutputWriter.getPartitionWriter(p)
       val outputTempFile = File.createTempFile("channelTemp", "", tempDir)
-      val outputTempFileStream = new FileOutputStream(outputTempFile)
-      outputTempFileStream.write(data(p))
-      outputTempFileStream.close()
+      Files.write(outputTempFile.toPath, data(p))
       val tempFileInput = new FileInputStream(outputTempFile)
       val channel = writer.openChannelWrapper()
       Utils.tryWithResource(new FileInputStream(outputTempFile)) { tempFileInput =>
-        Utils.tryWithResource(writer.openChannelWrapper()) { channelWrapper =>
+        Utils.tryWithResource(writer.openChannelWrapper().get) { channelWrapper =>
           assert(channelWrapper.channel().isInstanceOf[FileChannel],
             "Underlying channel should be a file channel")
           Utils.copyFileStreamNIO(