[SPARK-9716][ML] BinaryClassificationEvaluator should accept Double prediction column

This PR aims to allow the prediction column of `BinaryClassificationEvaluator` to be of double type.

Author: BenFradet <benjamin.fradet@gmail.com>

Closes #10472 from BenFradet/SPARK-9716.

1 files changed, 32 insertions, 0 deletions

diff --git a/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala
index a535c1218e..27349950dc 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/evaluation/BinaryClassificationEvaluatorSuite.scala
@@ -20,6 +20,7 @@ package org.apache.spark.ml.evaluation
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.ml.param.ParamsSuite
 import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.linalg.Vectors
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 
 class BinaryClassificationEvaluatorSuite
@@ -36,4 +37,35 @@ class BinaryClassificationEvaluatorSuite
       .setMetricName("areaUnderPR")
     testDefaultReadWrite(evaluator)
   }
+
+  test("should accept both vector and double raw prediction col") {
+    val evaluator = new BinaryClassificationEvaluator()
+      .setMetricName("areaUnderPR")
+
+    val vectorDF = sqlContext.createDataFrame(Seq(
+      (0d, Vectors.dense(12, 2.5)),
+      (1d, Vectors.dense(1, 3)),
+      (0d, Vectors.dense(10, 2))
+    )).toDF("label", "rawPrediction")
+    assert(evaluator.evaluate(vectorDF) === 1.0)
+
+    val doubleDF = sqlContext.createDataFrame(Seq(
+      (0d, 0d),
+      (1d, 1d),
+      (0d, 0d)
+    )).toDF("label", "rawPrediction")
+    assert(evaluator.evaluate(doubleDF) === 1.0)
+
+    val stringDF = sqlContext.createDataFrame(Seq(
+      (0d, "0d"),
+      (1d, "1d"),
+      (0d, "0d")
+    )).toDF("label", "rawPrediction")
+    val thrown = intercept[IllegalArgumentException] {
+      evaluator.evaluate(stringDF)
+    }
+    assert(thrown.getMessage.replace("\n", "") contains "Column rawPrediction must be of type " +
+      "equal to one of the following types: [DoubleType, ")
+    assert(thrown.getMessage.replace("\n", "") contains "but was actually of type StringType.")
+  }
 }