Refactor to remove code duplication and do the .get earlier

Author: a-roberts

core/src/main/scala/org/apache/spark/util/collection/PartitionedAppendOnlyMap.scala

@@ -30,25 +30,11 @@ private[spark] class PartitionedAppendOnlyMap[K, V]
 
   def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
     : Iterator[((Int, K), V)] = {
-    val comparator : Comparator[(Int, K)] =
-      if (keyComparator.isEmpty) {
-        partitionComparator
-      } else {
-        new Comparator[(Int, K)] {
-          override def compare(a: (Int, K), b: (Int, K)): Int = {
-          val partitionDiff = a._1 - b._1
-          if (partitionDiff != 0) {
-            partitionDiff
-          } else {
-            keyComparator.get.compare(a._2, b._2)
-          }
-        }
-      }
-    }
-    destructiveSortedIterator(comparator)
+    destructiveSortedIterator(getComparator(keyComparator))
   }
 
   def insert(partition: Int, key: K, value: V): Unit = {
     update((partition, key), value)
   }
 }
+

core/src/main/scala/org/apache/spark/util/collection/PartitionedPairBuffer.scala

@@ -74,22 +74,8 @@ private[spark] class PartitionedPairBuffer[K, V](initialCapacity: Int = 64)
   /** Iterate through the data in a given order. For this class this is not really destructive. */
   override def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
     : Iterator[((Int, K), V)] = {
-    val comparator : Comparator[(Int, K)] =
-      if (keyComparator.isEmpty) {
-        partitionComparator
-    } else {
-      new Comparator[(Int, K)] {
-        override def compare(a: (Int, K), b: (Int, K)): Int = {
-          val partitionDiff = a._1 - b._1
-            if (partitionDiff != 0) {
-              partitionDiff
-            } else {
-              keyComparator.get.compare(a._2, b._2)
-            }
-        }
-      }
-    }
-    new Sorter(new KVArraySortDataFormat[(Int, K), AnyRef]).sort(data, 0, curSize, comparator)
+    new Sorter(new KVArraySortDataFormat[(Int, K),
+      AnyRef]).sort(data, 0, curSize, getComparator(keyComparator))
     iterator
   }
 
@@ -112,3 +98,4 @@ private[spark] class PartitionedPairBuffer[K, V](initialCapacity: Int = 64)
 private object PartitionedPairBuffer {
   val MAXIMUM_CAPACITY = Int.MaxValue / 2 // 2 ^ 30 - 1
 }
+

core/src/main/scala/org/apache/spark/util/collection/WritablePartitionedPairCollection.scala

@@ -74,6 +74,31 @@ private[spark] object WritablePartitionedPairCollection {
     }
   }
 
+  /* Takes an optional parameter (keyComparator), use if provided
+   * and returns a comparator for the partitions
+   */
+  def getComparator[K](keyComparator: Option[Comparator[K]]) : Comparator[(Int, K)] = {
+    val comparator : Comparator[(Int, K)] =
+      if (keyComparator.isEmpty) {
+        partitionComparator
+      } else {
+        new Comparator[(Int, K)] {
+          // We know we have a non-empty comparator here
+          val ourKeyComp = keyComparator.get
+          override def compare(a: (Int, K), b: (Int, K)): Int = {
+            val partitionDiff = a._1 - b._1
+            if (partitionDiff != 0) {
+              partitionDiff
+            } else {
+              ourKeyComp.compare(a._2, b._2)
+            }
+          }
+        }
+      }
+    comparator
+  }
+
+
   /**
    * A comparator for (Int, K) pairs that orders them both by their partition ID and a key ordering.
    */
@@ -102,3 +127,4 @@ private[spark] trait WritablePartitionedIterator {
 
   def nextPartition(): Int
 }
+