callUDF => udf in Classifier where possible

maropu · Jun 24, 2015 · 0ea30b3 · 0ea30b3
1 parent 197ec82
commit 0ea30b3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala
@@ -108,7 +108,7 @@ abstract class ClassificationModel[FeaturesType, M <: ClassificationModel[Featur
     }
     if (getPredictionCol != "") {
       val predUDF = if (getRawPredictionCol != "") {
-        callUDF(raw2prediction _, DoubleType, col(getRawPredictionCol))
+        udf[Double, Vector](raw2prediction).apply(col(getRawPredictionCol))
       } else {
         callUDF(predict _, DoubleType, col(getFeaturesCol))
       }