SPARK[1784]: Adding a balancedPartitioner

bagel/src/test/scala/org/apache/spark/bagel/BagelSuite.scala

@@ -38,8 +38,6 @@ class BagelSuite extends FunSuite with Assertions with BeforeAndAfter with Timeo
       sc.stop()
       sc = null
     }
-    // To avoid Akka rebinding to the same port, since it doesn't unbind immediately on shutdown
-    System.clearProperty("spark.driver.port")
   }
 
   test("halting by voting") {

core/src/main/scala/org/apache/spark/Partitioner.scala

@@ -156,3 +156,41 @@ class RangePartitioner[K : Ordering : ClassTag, V](
       false
   }
 }
+
+/**
+ * A [[org.apache.spark.Partitioner]] that partitions records into balanced partitions
+ * by allocating keys to partitions in a round robin fashion.
+ */
+class BalancedPartitioner[K : Ordering : ClassTag, V](
+                                                    partitions: Int,
+                                                    @transient rdd: RDD[_ <: Product2[K,V]])
+  extends Partitioner {
+
+  // this array keeps track of keys assigned to a partition
+  // counts[0] refers to # of keys in partition 0 and so on
+  private val counts: Array[Int] = {
+    new Array[Int](numPartitions)
+  }
+
+  def numPartitions = math.abs(partitions)
+
+  var currPartition = 0
+
+  /*
+  * Pick current partition for the key in round robin manner
+  */
+  def getPartition(key: Any): Int = {
+    val partition = currPartition
+    counts(partition) = counts(partition) + 1
+    currPartition = (currPartition + 1) % numPartitions
+    partition
+  }
+
+  override def equals(other: Any): Boolean = other match {
+    case r: BalancedPartitioner[_,_] =>
+      (r.counts.sameElements(counts)
+        && r.currPartition == currPartition)
+    case _ =>
+      false
+  }
+}

core/src/test/java/org/apache/spark/JavaAPISuite.java

@@ -65,8 +65,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    // To avoid Akka rebinding to the same port, since it doesn't unbind immediately on shutdown
-    System.clearProperty("spark.driver.port");
   }
 
   static class ReverseIntComparator implements Comparator<Integer>, Serializable {

core/src/test/scala/org/apache/spark/AkkaUtilsSuite.scala

@@ -39,7 +39,6 @@ class AkkaUtilsSuite extends FunSuite with LocalSparkContext {
     val hostname = "localhost"
     val (actorSystem, boundPort) = AkkaUtils.createActorSystem("spark", hostname, 0,
       conf = conf, securityManager = securityManager)
-    System.setProperty("spark.driver.port", boundPort.toString)    // Will be cleared by LocalSparkContext
     System.setProperty("spark.hostPort", hostname + ":" + boundPort)
     assert(securityManager.isAuthenticationEnabled() === true)
 
@@ -77,7 +76,6 @@ class AkkaUtilsSuite extends FunSuite with LocalSparkContext {
     val hostname = "localhost"
     val (actorSystem, boundPort) = AkkaUtils.createActorSystem("spark", hostname, 0,
       conf = conf, securityManager = securityManager)
-    System.setProperty("spark.driver.port", boundPort.toString)    // Will be cleared by LocalSparkContext
     System.setProperty("spark.hostPort", hostname + ":" + boundPort)
 
     assert(securityManager.isAuthenticationEnabled() === false)
@@ -129,7 +127,6 @@ class AkkaUtilsSuite extends FunSuite with LocalSparkContext {
     val hostname = "localhost"
     val (actorSystem, boundPort) = AkkaUtils.createActorSystem("spark", hostname, 0,
       conf = conf, securityManager = securityManager)
-    System.setProperty("spark.driver.port", boundPort.toString)    // Will be cleared by LocalSparkContext
     System.setProperty("spark.hostPort", hostname + ":" + boundPort)
 
     assert(securityManager.isAuthenticationEnabled() === true)
@@ -182,7 +179,6 @@ class AkkaUtilsSuite extends FunSuite with LocalSparkContext {
     val hostname = "localhost"
     val (actorSystem, boundPort) = AkkaUtils.createActorSystem("spark", hostname, 0,
       conf = conf, securityManager = securityManager)
-    System.setProperty("spark.driver.port", boundPort.toString)    // Will be cleared by LocalSparkContext
     System.setProperty("spark.hostPort", hostname + ":" + boundPort)
 
     assert(securityManager.isAuthenticationEnabled() === true)

core/src/test/scala/org/apache/spark/MapOutputTrackerSuite.scala

@@ -124,9 +124,6 @@ class MapOutputTrackerSuite extends FunSuite with LocalSparkContext {
     val (actorSystem, boundPort) = AkkaUtils.createActorSystem("spark", hostname, 0, conf = conf,
       securityManager = new SecurityManager(conf))
 
-    // Will be cleared by LocalSparkContext
-    System.setProperty("spark.driver.port", boundPort.toString)
-
     val masterTracker = new MapOutputTrackerMaster(conf)
     masterTracker.trackerActor = actorSystem.actorOf(
       Props(new MapOutputTrackerMasterActor(masterTracker, conf)), "MapOutputTracker")

core/src/test/scala/org/apache/spark/PartitioningSuite.scala

@@ -66,6 +66,34 @@ class PartitioningSuite extends FunSuite with SharedSparkContext with PrivateMet
     assert(descendingP4 != p4)
   }
 
+  test("BalancedPartitioner equality") {
+    // Make an RDD where all the elements are the same so that the partition range bounds
+    // are deterministically all the same.
+    val rdd = sc.parallelize(1.to(4000)).map(x => (x, x))
+
+    val p2 = new BalancedPartitioner(2, rdd)
+    val p4 = new BalancedPartitioner(4, rdd)
+    val anotherP4 = new BalancedPartitioner(4, rdd)
+
+    assert(p2 === p2)
+    assert(p4 === p4)
+    assert(p2 != p4)
+    assert(p4 != p2)
+    assert(p4 === anotherP4)
+    assert(anotherP4 === p4)
+  }
+
+  test("BalancedPartitioner getPartition") {
+    val rdd = sc.parallelize(1.to(2000)).map(x => (x, x))
+    val partitioner = new BalancedPartitioner(4, rdd)
+    var expectedPartition = 0
+    1.to(2000).map { element => {
+      val partition = partitioner.getPartition(element)
+      assert(partition === expectedPartition)
+      expectedPartition = (expectedPartition + 1) % 4
+    }}
+  }
+
   test("RangePartitioner getPartition") {
     val rdd = sc.parallelize(1.to(2000)).map(x => (x, x))
     // We have different behaviour of getPartition for partitions with less than 1000 and more than

core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala

@@ -78,8 +78,6 @@ class BlockManagerSuite extends FunSuite with BeforeAndAfter with PrivateMethodT
   }
 
   after {
-    System.clearProperty("spark.driver.port")
-
     if (store != null) {
       store.stop()
       store = null

extras/java8-tests/src/test/java/org/apache/spark/Java8APISuite.java

@@ -57,8 +57,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    // To avoid Akka rebinding to the same port, since it doesn't unbind immediately on shutdown
-    System.clearProperty("spark.driver.port");
   }
 
   @Test

graphx/src/test/scala/org/apache/spark/graphx/LocalSparkContext.scala

@@ -38,8 +38,6 @@ trait LocalSparkContext {
       f(sc)
     } finally {
       sc.stop()
-      // To avoid Akka rebinding to the same port, since it doesn't unbind immediately on shutdown
-      System.clearProperty("spark.driver.port")
     }
   }
 }

mllib/src/test/java/org/apache/spark/mllib/classification/JavaLogisticRegressionSuite.java

@@ -42,7 +42,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    System.clearProperty("spark.driver.port");
   }
 
   int validatePrediction(List<LabeledPoint> validationData, LogisticRegressionModel model) {

mllib/src/test/java/org/apache/spark/mllib/classification/JavaNaiveBayesSuite.java

@@ -42,7 +42,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    System.clearProperty("spark.driver.port");
   }
 
   private static final List<LabeledPoint> POINTS = Arrays.asList(

mllib/src/test/java/org/apache/spark/mllib/classification/JavaSVMSuite.java

@@ -41,7 +41,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    System.clearProperty("spark.driver.port");
   }
 
   int validatePrediction(List<LabeledPoint> validationData, SVMModel model) {

mllib/src/test/java/org/apache/spark/mllib/clustering/JavaKMeansSuite.java

@@ -44,7 +44,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    System.clearProperty("spark.driver.port");
   }
 
   @Test

mllib/src/test/java/org/apache/spark/mllib/recommendation/JavaALSSuite.java

@@ -42,7 +42,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    System.clearProperty("spark.driver.port");
   }
 
   static void validatePrediction(MatrixFactorizationModel model, int users, int products, int features,

mllib/src/test/java/org/apache/spark/mllib/regression/JavaLassoSuite.java

@@ -41,7 +41,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    System.clearProperty("spark.driver.port");
   }
 
   int validatePrediction(List<LabeledPoint> validationData, LassoModel model) {

mllib/src/test/java/org/apache/spark/mllib/regression/JavaLinearRegressionSuite.java

@@ -41,7 +41,6 @@ public void setUp() {
   public void tearDown() {
     sc.stop();
     sc = null;
-    System.clearProperty("spark.driver.port");
   }
 
   int validatePrediction(List<LabeledPoint> validationData, LinearRegressionModel model) {

mllib/src/test/java/org/apache/spark/mllib/regression/JavaRidgeRegressionSuite.java

@@ -43,7 +43,6 @@ public void setUp() {
   public void tearDown() {
       sc.stop();
       sc = null;
-      System.clearProperty("spark.driver.port");
   }
 
   double predictionError(List<LabeledPoint> validationData, RidgeRegressionModel model) {

mllib/src/test/scala/org/apache/spark/mllib/util/LocalSparkContext.scala

@@ -34,7 +34,6 @@ trait LocalSparkContext extends BeforeAndAfterAll { self: Suite =>
     if (sc != null) {
       sc.stop()
     }
-    System.clearProperty("spark.driver.port")
     super.afterAll()
   }
 }

python/pyspark/tests.py

@@ -49,10 +49,6 @@ def setUp(self):
     def tearDown(self):
         self.sc.stop()
         sys.path = self._old_sys_path
-        # To avoid Akka rebinding to the same port, since it doesn't unbind
-        # immediately on shutdown
-        self.sc._jvm.System.clearProperty("spark.driver.port")
-
 
 class TestCheckpoint(PySparkTestCase):
 

repl/src/test/scala/org/apache/spark/repl/ReplSuite.scala

@@ -50,8 +50,6 @@ class ReplSuite extends FunSuite {
     if (interp.sparkContext != null) {
       interp.sparkContext.stop()
     }
-    // To avoid Akka rebinding to the same port, since it doesn't unbind immediately on shutdown
-    System.clearProperty("spark.driver.port")
     return out.toString
   }
 

sql/hive/src/main/scala/org/apache/spark/sql/hive/TestHive.scala

@@ -58,7 +58,6 @@ class TestHiveContext(sc: SparkContext) extends LocalHiveContext(sc) {
 
   // By clearing the port we force Spark to pick a new one.  This allows us to rerun tests
   // without restarting the JVM.
-  System.clearProperty("spark.driver.port")
   System.clearProperty("spark.hostPort")
 
   override lazy val warehousePath = getTempFilePath("sparkHiveWarehouse").getCanonicalPath

streaming/src/main/scala/org/apache/spark/streaming/util/MasterFailureTest.scala

@@ -186,7 +186,6 @@ object MasterFailureTest extends Logging {
     setupCalled = true
 
     // Setup the streaming computation with the given operation
-    System.clearProperty("spark.driver.port")
     val ssc = new StreamingContext("local[4]", "MasterFailureTest", batchDuration, null, Nil,
       Map())
     ssc.checkpoint(checkpointDir.toString)

streaming/src/test/java/org/apache/spark/streaming/LocalJavaStreamingContext.java

@@ -27,7 +27,6 @@ public abstract class LocalJavaStreamingContext {
 
     @Before
     public void setUp() {
-        System.clearProperty("spark.driver.port");
         System.setProperty("spark.streaming.clock", "org.apache.spark.streaming.util.ManualClock");
         ssc = new JavaStreamingContext("local[2]", "test", new Duration(1000));
         ssc.checkpoint("checkpoint");
@@ -37,8 +36,5 @@ public void setUp() {
     public void tearDown() {
         ssc.stop();
         ssc = null;
-
-        // To avoid Akka rebinding to the same port, since it doesn't unbind immediately on shutdown
-        System.clearProperty("spark.driver.port");
     }
 }

streaming/src/test/scala/org/apache/spark/streaming/CheckpointSuite.scala

@@ -368,7 +368,6 @@ class CheckpointSuite extends TestSuiteBase {
       "\n-------------------------------------------\n"
     )
     ssc = new StreamingContext(checkpointDir)
-    System.clearProperty("spark.driver.port")
     ssc.start()
     val outputNew = advanceTimeWithRealDelay[V](ssc, nextNumBatches)
     // the first element will be re-processed data of the last batch before restart

streaming/src/test/scala/org/apache/spark/streaming/TestSuiteBase.scala

@@ -153,8 +153,6 @@ trait TestSuiteBase extends FunSuite with BeforeAndAfter with Logging {
   // Default after function for any streaming test suite. Override this
   // if you want to add your stuff to "after" (i.e., don't call after { } )
   def afterFunction() {
-    // To avoid Akka rebinding to the same port, since it doesn't unbind immediately on shutdown
-    System.clearProperty("spark.driver.port")
     System.clearProperty("spark.streaming.clock")
   }