Adding NamedBroadcasts with persister

Author: carbribi

job-server-extras/src/spark.jobserver/NamedBroadcast.scala

@@ -0,0 +1,32 @@
+package spark.jobserver
+
+import org.apache.spark.broadcast.Broadcast
+
+/**
+ * wrapper for named objects of type Broadcast
+ */
+case class NamedBroadcast[T](broadcast: Broadcast[T]) extends NamedObject
+
+/**
+  * implementation of a NamedObjectPersister for Broadcast objects
+  */
+class BroadcastPersister[T] extends NamedObjectPersister[NamedBroadcast[T]] {
+  override def persist(namedObj: NamedBroadcast[T], name: String) {
+    namedObj match {
+      case NamedBroadcast(broadcast) =>{}
+    }
+  }
+  override def unpersist(namedObj: NamedBroadcast[T]) {
+    namedObj match {
+      case NamedBroadcast(broadcast) =>
+        broadcast.unpersist(blocking = false)
+    }
+  }
+  /**
+    * @param namedBroadcast the NamedBroadcast to refresh
+    */
+  override def refresh(namedBroadcast: NamedBroadcast[T]): NamedBroadcast[T] = namedBroadcast match {
+    case NamedBroadcast(broadcast) =>
+      namedBroadcast
+  }
+}

job-server-extras/src/spark.jobserver/NamedObjectsTestJob.scala

@@ -15,6 +15,7 @@ class NamedObjectsTestJob extends SparkJob with NamedObjectSupport {
   import NamedObjectsTestJobConfig._
   implicit def rddPersister: NamedObjectPersister[NamedRDD[Row]] = new RDDPersister[Row]
   implicit def dataFramePersister: NamedObjectPersister[NamedDataFrame] = new DataFramePersister
+  implicit def broadcastPersister[U]: NamedObjectPersister[NamedBroadcast[U]] = new BroadcastPersister[U]
 
   def validate(sql: SparkContext, config: Config): SparkJobValidation = SparkJobValid
 
@@ -35,6 +36,11 @@ class NamedObjectsTestJob extends SparkJob with NamedObjectSupport {
       namedObjects.update("rdd1", NamedRDD(rows(sc), true, StorageLevel.MEMORY_ONLY))
     }
 
+    if (config.hasPath(CREATE_BROADCAST)){
+      val broadcast = sc.broadcast(Set(1,2,3,4,5))
+      namedObjects.update("broadcast1", NamedBroadcast(broadcast))
+    }
+
     if (config.hasPath(DELETE)) {
       val iter = config.getStringList(DELETE).iterator
       while (iter.hasNext) {
@@ -49,6 +55,7 @@ class NamedObjectsTestJob extends SparkJob with NamedObjectSupport {
 object NamedObjectsTestJobConfig {
   val CREATE_DF = "createDF"
   val CREATE_RDD = "createRDD"
+  val CREATE_BROADCAST = "createBroadcast"
   val DELETE = "delete"
 }
 

job-server-extras/test/spark.jobserver/NamedObjectsJobSpec.scala

@@ -28,10 +28,11 @@ class NamedObjectsJobSpec extends JobSpecBase(JobManagerSpec.getNewSystem) {
 
   val jobName = "spark.jobserver.NamedObjectsTestJob"
 
-  private def getCreateConfig(createDF: Boolean, createRDD: Boolean) : Config = {
+  private def getCreateConfig(createDF: Boolean, createRDD: Boolean, createBroadcast: Boolean) : Config = {
     ConfigFactory.parseString("spark.jobserver.named-object-creation-timeout = 60 s, " + 
         NamedObjectsTestJobConfig.CREATE_DF + " = " + createDF + ", " +
-        NamedObjectsTestJobConfig.CREATE_RDD + " = " + createRDD)
+        NamedObjectsTestJobConfig.CREATE_RDD + " = " + createRDD + ", " +
+        NamedObjectsTestJobConfig.CREATE_BROADCAST + " = " + createBroadcast)
   }
 
   private def getDeleteConfig(names: List[String]) : Config = {
@@ -42,11 +43,11 @@ class NamedObjectsJobSpec extends JobSpecBase(JobManagerSpec.getNewSystem) {
   describe("NamedObjects (RDD)") {
     it("should survive from one job to another one") {
 
-      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, true), errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, true, false), errorEvents ++ syncEvents)
       val JobResult(_, names: Array[String]) = expectMsgClass(classOf[JobResult])
       names should contain("rdd1")
 
-      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false), errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false, false), errorEvents ++ syncEvents)
       val JobResult(_, names2: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names2 should contain("rdd1")
@@ -64,12 +65,12 @@ class NamedObjectsJobSpec extends JobSpecBase(JobManagerSpec.getNewSystem) {
   describe("NamedObjects (DataFrame)") {
     it("should survive from one job to another one") {
 
-      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(true, false), errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(true, false, false), errorEvents ++ syncEvents)
       val JobResult(_, names: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names should contain("df1")
 
-      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false), errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false, false), errorEvents ++ syncEvents)
       val JobResult(_, names2: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names2 should equal(names)
@@ -83,13 +84,13 @@ class NamedObjectsJobSpec extends JobSpecBase(JobManagerSpec.getNewSystem) {
   describe("NamedObjects (DataFrame + RDD)") {
     it("should survive from one job to another one") {
 
-      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(true, true), errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(true, true, false), errorEvents ++ syncEvents)
       val JobResult(_, names: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names should contain("rdd1")
       names should contain("df1")
 
-      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false), errorEvents ++ syncEvents)
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false, false), errorEvents ++ syncEvents)
       val JobResult(_, names2: Array[String]) = expectMsgClass(classOf[JobResult])
 
       names2 should equal(names)
@@ -100,4 +101,25 @@ class NamedObjectsJobSpec extends JobSpecBase(JobManagerSpec.getNewSystem) {
     }
   }
 
+  describe("NamedObjects (Broadcast)") {
+    it("should survive from one job to another one") {
+
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(true, true, true), errorEvents ++ syncEvents)
+      val JobResult(_, names: Array[String]) = expectMsgClass(classOf[JobResult])
+      
+      names should contain("rdd1")
+      names should contain("df1")
+      names should contain("broadcast1")
+
+      manager ! JobManagerActor.StartJob("demo", jobName, getCreateConfig(false, false, false), errorEvents ++ syncEvents)
+      val JobResult(_, names2: Array[String]) = expectMsgClass(classOf[JobResult])
+
+      names2 should equal(names)
+      
+      //clean-up
+      manager ! JobManagerActor.StartJob("demo", jobName, getDeleteConfig(List("rdd1", "df1", "broadcast1"))
+        , errorEvents ++ syncEvents)
+      val JobResult(_, names3: Array[String]) = expectMsgClass(classOf[JobResult])
+    }
+  }
 }