Merge pull request #115 from aarondav/shuffle-fix

Eliminate extra memory usage when shuffle file consolidation is disabled

Otherwise, we see SPARK-946 even when shuffle file consolidation is disabled.
Fixing SPARK-946 is still forthcoming.

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -523,7 +523,9 @@ private[spark] class BlockManager(
     val file = diskBlockManager.createBlockFile(blockId, filename, allowAppending =  true)
     val writer = new DiskBlockObjectWriter(blockId, file, serializer, bufferSize, compressStream)
     writer.registerCloseEventHandler(() => {
-      diskBlockManager.mapBlockToFileSegment(blockId, writer.fileSegment())
+      if (shuffleBlockManager.consolidateShuffleFiles) {
+        diskBlockManager.mapBlockToFileSegment(blockId, writer.fileSegment())
+      }
       val myInfo = new BlockInfo(StorageLevel.DISK_ONLY, false)
       blockInfo.put(blockId, myInfo)
       myInfo.markReady(writer.fileSegment().length)

core/src/main/scala/org/apache/spark/storage/ShuffleBlockManager.scala

@@ -63,8 +63,12 @@ class ShuffleBlockManager(blockManager: BlockManager) {
         val fileId = getUnusedFileId()
         val writers = Array.tabulate[BlockObjectWriter](numBuckets) { bucketId =>
           val blockId = ShuffleBlockId(shuffleId, mapId, bucketId)
-          val filename = physicalFileName(shuffleId, bucketId, fileId)
-          blockManager.getDiskWriter(blockId, filename, serializer, bufferSize)
+          if (consolidateShuffleFiles) {
+            val filename = physicalFileName(shuffleId, bucketId, fileId)
+            blockManager.getDiskWriter(blockId, filename, serializer, bufferSize)
+          } else {
+            blockManager.getDiskWriter(blockId, blockId.name, serializer, bufferSize)
+          }
         }
         new ShuffleWriterGroup(mapId, fileId, writers)
       }
@@ -81,8 +85,9 @@ class ShuffleBlockManager(blockManager: BlockManager) {
   }
 
   private def recycleFileId(fileId: Int) {
-    if (!consolidateShuffleFiles) { return } // ensures we always generate new file id
-    unusedFileIds.add(fileId)
+    if (consolidateShuffleFiles) {
+      unusedFileIds.add(fileId)
+    }
   }
 
   private def physicalFileName(shuffleId: Int, bucketId: Int, fileId: Int) = {