From 85da63fb741b22a21b147b28aa4145e884f91a50 Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Tue, 12 May 2015 12:09:09 -0700
Subject: [PATCH] Cleanup in UnsafeShuffleSorterIterator.

---
 .../unsafe/UnsafeShuffleInMemorySorter.java   | 37 ++++++++++---------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java b/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java
index 6fb87848df38b..57b24547125f6 100644
--- a/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java
@@ -78,14 +78,29 @@ public void insertRecord(long recordPointer, int partitionId) {
     sortBufferInsertPosition++;
   }
 
-  public static abstract class UnsafeShuffleSorterIterator {
+  /**
+   * An iterator-like class that's used instead of Java's Iterator in order to facilitate inlining.
+   */
+  public static final class UnsafeShuffleSorterIterator {
 
+    private final long[] sortBuffer;
+    private final int numRecords;
     final PackedRecordPointer packedRecordPointer = new PackedRecordPointer();
+    private int position = 0;
 
-    public abstract boolean hasNext();
+    public UnsafeShuffleSorterIterator(int numRecords, long[] sortBuffer) {
+      this.numRecords = numRecords;
+      this.sortBuffer = sortBuffer;
+    }
 
-    public abstract void loadNext();
+    public boolean hasNext() {
+      return position < numRecords;
+    }
 
+    public void loadNext() {
+      packedRecordPointer.set(sortBuffer[position]);
+      position++;
+    }
   }
 
   /**
@@ -93,20 +108,6 @@ public static abstract class UnsafeShuffleSorterIterator {
    */
   public UnsafeShuffleSorterIterator getSortedIterator() {
     sorter.sort(sortBuffer, 0, sortBufferInsertPosition, SORT_COMPARATOR);
-    return new UnsafeShuffleSorterIterator() {
-
-      private int position = 0;
-
-      @Override
-      public boolean hasNext() {
-        return position < sortBufferInsertPosition;
-      }
-
-      @Override
-      public void loadNext() {
-        packedRecordPointer.set(sortBuffer[position]);
-        position++;
-      }
-    };
+    return new UnsafeShuffleSorterIterator(sortBufferInsertPosition, sortBuffer);
   }
 }