[SPARK-7127] [MLLIB] Adding broadcast of model before prediction for ensembles

mllib/src/main/scala/org/apache/spark/ml/Predictor.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.ml
 
 import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.ml.param._
 import org.apache.spark.ml.param.shared._
 import org.apache.spark.ml.util.SchemaUtils
@@ -176,17 +177,53 @@ abstract class PredictionModel[FeaturesType, M <: PredictionModel[FeaturesType,
   override def transform(dataset: DataFrame): DataFrame = {
     transformSchema(dataset.schema, logging = true)
     if ($(predictionCol).nonEmpty) {
-      dataset.withColumn($(predictionCol), callUDF(predict _, DoubleType, col($(featuresCol))))
+      transformImpl(dataset)
     } else {
       this.logWarning(s"$uid: Predictor.transform() was called as NOOP" +
         " since no output columns were set.")
       dataset
     }
   }
 
+  protected def transformImpl(dataset: DataFrame): DataFrame = {
+    dataset.withColumn($(predictionCol), callUDF(predict _, DoubleType, col($(featuresCol))))
+  }
+
   /**
    * Predict label for the given features.
    * This internal method is used to implement [[transform()]] and output [[predictionCol]].
    */
   protected def predict(features: FeaturesType): Double
 }
+
+
+/**
+ * :: DeveloperApi ::
+ *
+ * Abstraction for a model for prediction tasks that will broadcast the model used to predict.
+ *
+ * @tparam FeaturesType  Type of features.
+ *                       E.g., [[org.apache.spark.mllib.linalg.VectorUDT]] for vector features.
+ * @tparam M  Specialization of [[PredictionModel]].  If you subclass this type, use this type
+ *            parameter to specify the concrete type for the corresponding model.
+ */
+@DeveloperApi
+abstract class PredictionModelBroadcasting[
+    FeaturesType, M <: PredictionModelBroadcasting[FeaturesType, M]
+  ]
+  extends PredictionModel[FeaturesType, M] {
+
+  protected def transformImpl(dataset: DataFrame, bcastModel: Broadcast[M]): DataFrame = {
+
+    dataset.withColumn($(predictionCol),
+      callUDF((features: FeaturesType) => predictWithBroadcastModel(features, bcastModel),
+        DoubleType, col($(featuresCol)))
+    )
+  }
+
+  /**
+   * Predict label for the given features using a broadcasted model.
+   * This internal method is used to implement [[transform()]] and output [[predictionCol]].
+   */
+  protected def predictWithBroadcastModel(features: FeaturesType, bcastModel: Broadcast[M]): Double
+}

mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala

@@ -20,7 +20,8 @@ package org.apache.spark.ml.classification
 import scala.collection.mutable
 
 import org.apache.spark.annotation.AlphaComponent
-import org.apache.spark.ml.{PredictionModel, Predictor}
+import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.ml.{PredictionModelBroadcasting, Predictor}
 import org.apache.spark.ml.param.ParamMap
 import org.apache.spark.ml.tree.{RandomForestParams, TreeClassifierParams, DecisionTreeModel, TreeEnsembleModel}
 import org.apache.spark.ml.util.{Identifiable, MetadataUtils}
@@ -109,6 +110,7 @@ object RandomForestClassifier {
     RandomForestParams.supportedFeatureSubsetStrategies
 }
 
+
 /**
  * :: AlphaComponent ::
  *
@@ -122,7 +124,7 @@ object RandomForestClassifier {
 final class RandomForestClassificationModel private[ml] (
     override val uid: String,
     private val _trees: Array[DecisionTreeClassificationModel])
-  extends PredictionModel[Vector, RandomForestClassificationModel]
+  extends PredictionModelBroadcasting[Vector, RandomForestClassificationModel]
   with TreeEnsembleModel with Serializable {
 
   require(numTrees > 0, "RandomForestClassificationModel requires at least 1 tree.")
@@ -134,13 +136,28 @@ final class RandomForestClassificationModel private[ml] (
 
   override def treeWeights: Array[Double] = _treeWeights
 
+  override def transform(dataset: DataFrame): DataFrame = {
+    val bcastModel = dataset.sqlContext.sparkContext.broadcast(this)
+    transformImpl(dataset, bcastModel)
+  }
+
   override protected def predict(features: Vector): Double = {
-    // TODO: Override transform() to broadcast model.  SPARK-7127
     // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
+    // Predict without using a broadcasted mode
+    predictImpl(features, () => this)
+  }
+
+  override protected def predictWithBroadcastModel(features: Vector,
+      bcastModel: Broadcast[RandomForestClassificationModel]): Double = {
+    // Predict using the given broadcasted model
+    predictImpl(features, () => bcastModel.value)
+  }
+
+  protected def predictImpl(features: Vector, modelAccesor: () => TreeEnsembleModel): Double = {
     // Classifies using majority votes.
     // Ignore the weights since all are 1.0 for now.
     val votes = mutable.Map.empty[Int, Double]
-    _trees.view.foreach { tree =>
+    modelAccesor().trees.view.foreach { tree =>
       val prediction = tree.rootNode.predict(features).toInt
       votes(prediction) = votes.getOrElse(prediction, 0.0) + 1.0 // 1.0 = weight
     }