Review feedback.

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

@@ -340,6 +340,9 @@ private long[] mergeSpillsWithFileStream(
     final int numPartitions = partitioner.numPartitions();
     final long[] partitionLengths = new long[numPartitions];
     final InputStream[] spillInputStreams = new FileInputStream[spills.length];
+
+    // Use a counting output stream to avoid having to close the underlying file and ask
+    // the file system for its size after each partition is written.
     final CountingOutputStream mergedFileOutputStream = new CountingOutputStream(
       new FileOutputStream(outputFile));
 
@@ -350,6 +353,8 @@ private long[] mergeSpillsWithFileStream(
       }
       for (int partition = 0; partition < numPartitions; partition++) {
         final long initialFileLength = mergedFileOutputStream.getByteCount();
+        // Shield the underlying output stream from close() calls, so that we can close the higher
+        // level streams to make sure all data is really flushed and internal state is cleaned.
         OutputStream partitionOutput = new CloseShieldOutputStream(mergedFileOutputStream);
         partitionOutput = blockManager.serializerManager().wrapForEncryption(partitionOutput);
         if (compressionCodec != null) {
@@ -361,12 +366,16 @@ private long[] mergeSpillsWithFileStream(
           if (partitionLengthInSpill > 0) {
             InputStream partitionInputStream = new LimitedInputStream(spillInputStreams[i],
               partitionLengthInSpill, false);
-            partitionInputStream = blockManager.serializerManager().wrapForEncryption(
-              partitionInputStream);
-            if (compressionCodec != null) {
-              partitionInputStream = compressionCodec.compressedInputStream(partitionInputStream);
+            try {
+              partitionInputStream = blockManager.serializerManager().wrapForEncryption(
+                partitionInputStream);
+              if (compressionCodec != null) {
+                partitionInputStream = compressionCodec.compressedInputStream(partitionInputStream);
+              }
+              ByteStreams.copy(partitionInputStream, partitionOutput);
+            } finally {
+              partitionInputStream.close();
             }
-            ByteStreams.copy(partitionInputStream, partitionOutput);
           }
         }
         partitionOutput.flush();

core/src/main/scala/org/apache/spark/serializer/SerializerManager.scala

@@ -144,14 +144,14 @@ private[spark] class SerializerManager(
   /**
    * Wrap an output stream for compression if block compression is enabled for its block type
    */
-  def wrapForCompression(blockId: BlockId, s: OutputStream): OutputStream = {
+  private[this] def wrapForCompression(blockId: BlockId, s: OutputStream): OutputStream = {
     if (shouldCompress(blockId)) compressionCodec.compressedOutputStream(s) else s
   }
 
   /**
    * Wrap an input stream for compression if block compression is enabled for its block type
    */
-  def wrapForCompression(blockId: BlockId, s: InputStream): InputStream = {
+  private[this] def wrapForCompression(blockId: BlockId, s: InputStream): InputStream = {
     if (shouldCompress(blockId)) compressionCodec.compressedInputStream(s) else s
   }
 

core/src/test/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriterSuite.java

@@ -410,6 +410,18 @@ public void mergeSpillsWithCompressionAndEncryptionSlowPath() throws Exception {
     testMergingSpills(false, LZ4CompressionCodec.class.getName(), true);
   }
 
+  @Test
+  public void mergeSpillsWithEncryptionAndNoCompression() throws Exception {
+    // This should actually be translated to a "file stream merge" internally, just have the
+    // test to make sure that it's the case.
+    testMergingSpills(true, null, true);
+  }
+
+  @Test
+  public void mergeSpillsWithFileStreamAndEncryptionAndNoCompression() throws Exception {
+    testMergingSpills(false, null, true);
+  }
+
   @Test
   public void writeEnoughDataToTriggerSpill() throws Exception {
     memoryManager.limit(PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES);

core/src/test/scala/org/apache/spark/util/collection/ExternalAppendOnlyMapSuite.scala

@@ -17,12 +17,9 @@
 
 package org.apache.spark.util.collection
 
-import java.security.PrivilegedExceptionAction
-
 import scala.collection.mutable.ArrayBuffer
 
 import org.apache.spark._
-import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.internal.config._
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.memory.MemoryTestingUtils