Refactoring

Author: a-roberts

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

@@ -74,8 +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)] = {
-    new Sorter(new KVArraySortDataFormat[(Int, K),
-      AnyRef]).sort(data, 0, curSize, getComparator(keyComparator))
+    val comparator = getComparator(keyComparator)
+    new Sorter(new KVArraySortDataFormat[(Int, K), AnyRef]).sort(data, 0, curSize, comparator)
     iterator
   }
 

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

@@ -74,44 +74,25 @@ private[spark] object WritablePartitionedPairCollection {
     }
   }
 
-  /* Takes an optional parameter (keyComparator), use if provided
+  /**
+   * 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)
-            }
+  def getComparator[K](keyComparator: Option[Comparator[K]]): Comparator[(Int, K)] = {
+    if (!keyComparator.isDefined) return partitionComparator
+    else {
+      val theKeyComp = keyComparator.get
+      new Comparator[(Int, K)] {
+        // We know we have a non-empty comparator here
+        override def compare(a: (Int, K), b: (Int, K)): Int = {
+          val partitionDiff = a._1 - b._1
+          if (partitionDiff != 0) {
+            partitionDiff
+          } else {
+            theKeyComp.compare(a._2, b._2)
           }
         }
       }
-    comparator
-  }
-
-
-  /**
-   * A comparator for (Int, K) pairs that orders them both by their partition ID and a key ordering.
-   */
-  def partitionKeyComparator[K](keyComparator: Comparator[K]): Comparator[(Int, K)] = {
-    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.compare(a._2, b._2)
-        }
-      }
     }
   }
 }