Takes locality and the sum of partition length into account.

apache · Jul 20, 2015 · cb72d0f · cb72d0f
1 parent 79ec072
commit cb72d0f
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/core/src/main/scala/org/apache/spark/rdd/CoalescedRDD.scala b/core/src/main/scala/org/apache/spark/rdd/CoalescedRDD.scala
@@ -295,7 +295,15 @@ private class PartitionCoalescer(maxPartitions: Int, prev: RDD[_], balanceSlack:
 
     val r1 = rnd.nextInt(groupArr.size)
     val r2 = rnd.nextInt(groupArr.size)
-    val minPowerOfTwo = if (groupArr(r1).size < groupArr(r2).size) groupArr(r1) else groupArr(r2)
+    val minPowerOfTwo = if (p.isInstanceOf[HadoopPartition]) {
+      val groupLen1 = groupArr(r1).arr.map(part =>
+        part.asInstanceOf[HadoopPartition].inputSplit.value.getLength).sum
+      val groupLen2 = groupArr(r1).arr.map(part =>
+        part.asInstanceOf[HadoopPartition].inputSplit.value.getLength).sum
+      if (groupLen1 < groupLen2) groupArr(r1) else groupArr(r2)
+    } else {
+      if (groupArr(r1).size < groupArr(r2).size) groupArr(r1) else groupArr(r2)
+    }
     if (prefPart.isEmpty) {
       // if no preferred locations, just use basic power of two
       return minPowerOfTwo