address comments

Author: cloud-fan

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

@@ -422,11 +422,14 @@ private long[] mergeSpillsWithTransferTo(SpillInfo[] spills, File outputFile) th
       for (int partition = 0; partition < numPartitions; partition++) {
         for (int i = 0; i < spills.length; i++) {
           final long partitionLengthInSpill = spills[i].partitionLengths[partition];
-          long bytesToTransfer = partitionLengthInSpill;
           final FileChannel spillInputChannel = spillInputChannels[i];
           final long writeStartTime = System.nanoTime();
-          Utils.copyFileStreamNIO(spillInputChannel, mergedFileOutputChannel, bytesToTransfer);
-          spillInputChannelPositions[i] += bytesToTransfer;
+          Utils.copyFileStreamNIO(
+            spillInputChannel,
+            mergedFileOutputChannel,
+            spillInputChannelPositions[i],
+            partitionLengthInSpill);
+          spillInputChannelPositions[i] += partitionLengthInSpill;
           writeMetrics.incWriteTime(System.nanoTime() - writeStartTime);
           bytesWrittenToMergedFile += partitionLengthInSpill;
           partitionLengths[partition] += partitionLengthInSpill;

core/src/main/scala/org/apache/spark/util/Utils.scala

@@ -330,7 +330,7 @@ private[spark] object Utils extends Logging {
         val inChannel = in.asInstanceOf[FileInputStream].getChannel()
         val outChannel = out.asInstanceOf[FileOutputStream].getChannel()
         val size = inChannel.size()
-        copyFileStreamNIO(inChannel, outChannel, size)
+        copyFileStreamNIO(inChannel, outChannel, 0, size)
         size
       } else {
         var count = 0L
@@ -356,13 +356,19 @@ private[spark] object Utils extends Logging {
     }
   }
 
-  def copyFileStreamNIO(input: FileChannel, output: FileChannel, bytesToCopy: Long): Unit = {
+  def copyFileStreamNIO(
+      input: FileChannel,
+      output: FileChannel,
+      startPosition: Long,
+      bytesToCopy: Long): Unit = {
     val initialPos = output.position()
     var count = 0L
     // In case transferTo method transferred less data than we have required.
     while (count < bytesToCopy) {
-      count += input.transferTo(count, bytesToCopy - count, output)
+      count += input.transferTo(count + startPosition, bytesToCopy - count, output)
     }
+    assert(count == bytesToCopy,
+      s"request to copy $bytesToCopy bytes, but actually copied $count bytes.")
 
     // Check the position after transferTo loop to see if it is in the right position and
     // give user information if not.