Fix bug in calculating free space in current page.

This broke off-heap mode.
nemccarthy · May 3, 2015 · 8958584 · 8958584
1 parent f17fa8f
commit 8958584
Showing 1 changed file with 8 additions and 2 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
@@ -71,7 +71,7 @@ public class UnsafeShuffleWriter<K, V> implements ShuffleWriter<K, V> {
   private MapStatus mapStatus = null;
 
   private MemoryBlock currentPage = null;
-  private long currentPagePosition = PAGE_SIZE;
+  private long currentPagePosition = -1;
 
   /**
    * Are we in the process of stopping? Because map tasks can call stop() with success = true
@@ -110,11 +110,17 @@ public void write(scala.collection.Iterator<Product2<K, V>> records) {
   }
 
   private void ensureSpaceInDataPage(long requiredSpace) throws Exception {
+    final long spaceInCurrentPage;
+    if (currentPage != null) {
+      spaceInCurrentPage = PAGE_SIZE - (currentPagePosition - currentPage.getBaseOffset());
+    } else {
+      spaceInCurrentPage = 0;
+    }
     if (requiredSpace > PAGE_SIZE) {
       // TODO: throw a more specific exception?
       throw new Exception("Required space " + requiredSpace + " is greater than page size (" +
         PAGE_SIZE + ")");
-    } else if (requiredSpace > (PAGE_SIZE - currentPagePosition)) {
+    } else if (requiredSpace > spaceInCurrentPage) {
       currentPage = memoryManager.allocatePage(PAGE_SIZE);
       currentPagePosition = currentPage.getBaseOffset();
       allocatedPages.add(currentPage);