[SPARK-10017] [MLlib]: ML model broadcasts should be stored in private vars: mllib NaiveBayes

mllib/src/main/scala/org/apache/spark/mllib/classification/NaiveBayes.scala

@@ -21,16 +21,15 @@ import java.lang.{Iterable => JIterable}
 
 import scala.collection.JavaConverters._
 
-import org.json4s.JsonDSL._
-import org.json4s.jackson.JsonMethods._
-
 import org.apache.spark.{Logging, SparkContext, SparkException}
 import org.apache.spark.annotation.Since
 import org.apache.spark.mllib.linalg.{BLAS, DenseMatrix, DenseVector, SparseVector, Vector}
 import org.apache.spark.mllib.regression.LabeledPoint
-import org.apache.spark.mllib.util.{Loader, Saveable}
+import org.apache.spark.mllib.util.{Broadcastable, Loader, Saveable}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{DataFrame, SQLContext}
+import org.json4s.JsonDSL._
+import org.json4s.jackson.JsonMethods._
 
 /**
  * Model for Naive Bayes Classifiers.
@@ -47,7 +46,7 @@ class NaiveBayesModel private[spark] (
     @Since("0.9.0") val pi: Array[Double],
     @Since("0.9.0") val theta: Array[Array[Double]],
     @Since("1.4.0") val modelType: String)
-  extends ClassificationModel with Serializable with Saveable {
+  extends ClassificationModel with Serializable with Saveable with Broadcastable[NaiveBayesModel] {
 
   import NaiveBayes.{Bernoulli, Multinomial, supportedModelTypes}
 
@@ -86,9 +85,10 @@ class NaiveBayesModel private[spark] (
 
   @Since("1.0.0")
   override def predict(testData: RDD[Vector]): RDD[Double] = {
-    val bcModel = testData.context.broadcast(this)
+    val sc = testData.sparkContext
+    val lclBcModel = getBroadcastModel(sc, this)
     testData.mapPartitions { iter =>
-      val model = bcModel.value
+      val model = lclBcModel.value
       iter.map(model.predict)
     }
   }
@@ -112,9 +112,10 @@ class NaiveBayesModel private[spark] (
    */
   @Since("1.5.0")
   def predictProbabilities(testData: RDD[Vector]): RDD[Vector] = {
-    val bcModel = testData.context.broadcast(this)
+    val sc = testData.sparkContext
+    val lclBcModel = getBroadcastModel(sc, this)
     testData.mapPartitions { iter =>
-      val model = bcModel.value
+      val model = lclBcModel.value
       iter.map(model.predictProbabilities)
     }
   }

mllib/src/main/scala/org/apache/spark/mllib/util/ModelBroadcast.scala

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.mllib.util
+
+import scala.reflect.ClassTag
+
+import org.apache.spark.SparkContext
+import org.apache.spark.broadcast.Broadcast
+
+private[spark] trait Broadcastable[T] {
+
+  private var bcModel: Option[Broadcast[T]] = None
+
+  /**
+   * Checks whether the model object is already broadcast and returns the reference.
+   * If not, then broadcasts the model and returns a reference
+   * @param sc SparkContext that will be used for the broadcast
+   * @param modelToBc Model object to broadcast
+   * @return the broadcast model
+   */
+  def getBroadcastModel(sc: SparkContext, modelToBc: T)
+                       (implicit ev: ClassTag[T]) : Broadcast[T] = {
+    bcModel match {
+      case None => bcModel = Some(sc.broadcast(modelToBc))
+      case _ =>
+    }
+    bcModel.get
+  }
+
+}