From f6f7ca9d2ef65da15f42085993e58e618637fad5 Mon Sep 17 00:00:00 2001 From: BenFradet Date: Tue, 19 Jan 2016 14:59:20 -0800 Subject: [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 Closes #10472 from BenFradet/SPARK-9716. --- python/pyspark/ml/evaluation.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) (limited to 'python/pyspark/ml/evaluation.py') diff --git a/python/pyspark/ml/evaluation.py b/python/pyspark/ml/evaluation.py index dcc1738ec5..6ff68abd8f 100644 --- a/python/pyspark/ml/evaluation.py +++ b/python/pyspark/ml/evaluation.py @@ -106,8 +106,9 @@ class JavaEvaluator(Evaluator, JavaWrapper): @inherit_doc class BinaryClassificationEvaluator(JavaEvaluator, HasLabelCol, HasRawPredictionCol): """ - Evaluator for binary classification, which expects two input - columns: rawPrediction and label. + Evaluator for binary classification, which expects two input columns: rawPrediction and label. + The rawPrediction column can be of type double (binary 0/1 prediction, or probability of label + 1) or of type vector (length-2 vector of raw predictions, scores, or label probabilities). >>> from pyspark.mllib.linalg import Vectors >>> scoreAndLabels = map(lambda x: (Vectors.dense([1.0 - x[0], x[0]]), x[1]), -- cgit v1.2.3