Minor fix to cleanup logic.

jeanlyn · May 28, 2015 · 03f35a4 · 03f35a4
1 parent b5cc35b
commit 03f35a4
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java
@@ -120,6 +120,10 @@ public void insertAll(Iterator<Product2<K, V>> records) throws IOException {
       final K key = record._1();
       partitionWriters[partitioner.getPartition(key)].write(key, record._2());
     }
+
+    for (BlockObjectWriter writer : partitionWriters) {
+      writer.commitAndClose();
+    }
   }
 
   @Override
@@ -133,9 +137,7 @@ public long[] writePartitionedFile(
       // We were passed an empty iterator
       return lengths;
     }
-    for (BlockObjectWriter writer : partitionWriters) {
-      writer.commitAndClose();
-    }
+
     final FileOutputStream out = new FileOutputStream(outputFile, true);
     final long writeStartTime = System.nanoTime();
     boolean threwException = true;
@@ -158,6 +160,7 @@ public long[] writePartitionedFile(
       Closeables.close(out, threwException);
       writeMetrics.incShuffleWriteTime(System.nanoTime() - writeStartTime);
     }
+    partitionWriters = null;
     return lengths;
   }