Spark partitioning test with multiple RDDs

PartitioningTest.scala

import java.lang.management.ManagementFactory
import java.net.InetAddress
import org.apache.spark.rdd.RDD
import org.apache.spark.{Partitioner, SparkContext}
import scala.runtime.ScalaRunTime

/**
  * Note: Package as a jar and run with spark-submit against a running cluster.
  * Created by bekce on 6/5/17.
  */
object PartitioningTest {

  class OneDimPartitioner(size: Int, perPart: Int) extends Partitioner {
    private val parts = math.ceil(size * 1.0 / perPart).toInt
    override def numPartitions: Int = parts
    override def getPartition(key: Any): Int = key.asInstanceOf[Int] / perPart
  }

  def main(args: Array[String]): Unit = {
    val sc = new SparkContext()
    println(s"CONF=${ScalaRunTime.stringOf(sc.getConf.getAll)}, sc.defaultParallelism=${sc.defaultParallelism}")
    
    val partitioner = new OneDimPartitioner(120, 6)
    val A = sc.parallelize(0 to 119).map(t => (t, "A"+t)).partitionBy(partitioner).cache()
    printAll(A, "A")
    val B = sc.parallelize(0 to 119).map(t => (t, "B"+t)).partitionBy(partitioner).cache()
    printAll(B, "B")
    val C: RDD[(Int, (String, String))] = A.join(B, partitioner)
    printAll(C, "C")

    sc.stop()
  }

  def printAll(rdd: RDD[_ <: AnyRef], msg: String) : Unit = {
    rdd.foreach(t => {
      println(s"$msg, $t, jvm=${ManagementFactory.getRuntimeMXBean().getName()}, localhost=${InetAddress.getLocalHost()}")
    })
  }
}