Re-use a buffer in UnsafeShuffleWriter

jeanlyn · May 2, 2015 · 026b497 · 026b497
1 parent 0748458
commit 026b497
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleWriter.java b/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleWriter.java
@@ -193,6 +193,7 @@ private long[] writeSortedRecordsToFile(
     int currentPartition = -1;
     BlockObjectWriter writer = null;
 
+    final byte[] arr = new byte[SER_BUFFER_SIZE];
     while (sortedRecords.hasNext()) {
       final RecordPointerAndKeyPrefix recordPointer = sortedRecords.next();
       final int partition = (int) recordPointer.keyPrefix;
@@ -211,16 +212,14 @@ private long[] writeSortedRecordsToFile(
       final Object baseObject = memoryManager.getPage(recordPointer.recordPointer);
       final long baseOffset = memoryManager.getOffsetInPage(recordPointer.recordPointer);
       final int recordLength = (int) PlatformDependent.UNSAFE.getLong(baseObject, baseOffset + 8);
-      // TODO: re-use a buffer or avoid double-buffering entirely
-      final byte[] arr = new byte[recordLength];
-        PlatformDependent.copyMemory(
-          baseObject,
-          baseOffset + 16,
-          arr,
-          PlatformDependent.BYTE_ARRAY_OFFSET,
-          recordLength);
+      PlatformDependent.copyMemory(
+        baseObject,
+        baseOffset + 16,
+        arr,
+        PlatformDependent.BYTE_ARRAY_OFFSET,
+        recordLength);
       assert (writer != null);  // To suppress an IntelliJ warning
-      writer.write(arr);
+      writer.write(arr, 0, recordLength);
       // TODO: add a test that detects whether we leave this call out:
       writer.recordWritten();
     }