aboutsummaryrefslogtreecommitdiff
path: root/python/pyspark/ml/evaluation.py
diff options
context:
space:
mode:
authorBenFradet <benjamin.fradet@gmail.com>2016-01-19 14:59:20 -0800
committerJoseph K. Bradley <joseph@databricks.com>2016-01-19 14:59:20 -0800
commitf6f7ca9d2ef65da15f42085993e58e618637fad5 (patch)
tree5c2d266d6b48111ff853ea4bb468c08038e3c0fa /python/pyspark/ml/evaluation.py
parent43f1d59e17d89d19b322d639c5069a3fc0c8e2ed (diff)
downloadspark-f6f7ca9d2ef65da15f42085993e58e618637fad5.tar.gz
spark-f6f7ca9d2ef65da15f42085993e58e618637fad5.tar.bz2
spark-f6f7ca9d2ef65da15f42085993e58e618637fad5.zip
[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.
Diffstat (limited to 'python/pyspark/ml/evaluation.py')
-rw-r--r--python/pyspark/ml/evaluation.py5
1 files changed, 3 insertions, 2 deletions
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]),