From dfc4c935ba99a4adab2bf39cf24fff1ee3011759 Mon Sep 17 00:00:00 2001
From: Yanbo Liang <ybliang8@gmail.com>
Date: Thu, 5 Jan 2017 18:59:49 -0800
Subject: [PATCH] [MINOR] Correct LogisticRegression test case for
probability2prediction.
## What changes were proposed in this pull request?
Set correct column names for ```force to use probability2prediction``` in ```LogisticRegressionSuite```.
## How was this patch tested?
Change unit test.
Author: Yanbo Liang <ybliang8@gmail.com>
Closes #16477 from yanboliang/lor-pred.
---
.../spark/ml/classification/LogisticRegressionSuite.scala | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
index 1308210417..c14dcbd552 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/LogisticRegressionSuite.scala
@@ -360,7 +360,7 @@ class LogisticRegressionSuite
}
// force it to use raw2prediction
- model.setProbabilityCol("")
+ model.setRawPredictionCol("rawPrediction").setProbabilityCol("")
val resultsUsingRaw2Predict =
model.transform(smallMultinomialDataset).select("prediction").as[Double].collect()
resultsUsingRaw2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
@@ -368,7 +368,7 @@ class LogisticRegressionSuite
}
// force it to use probability2prediction
- model.setRawPredictionCol("")
+ model.setRawPredictionCol("").setProbabilityCol("probability")
val resultsUsingProb2Predict =
model.transform(smallMultinomialDataset).select("prediction").as[Double].collect()
resultsUsingProb2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
@@ -414,7 +414,7 @@ class LogisticRegressionSuite
}
// force it to use raw2prediction
- model.setProbabilityCol("")
+ model.setRawPredictionCol("rawPrediction").setProbabilityCol("")
val resultsUsingRaw2Predict =
model.transform(smallBinaryDataset).select("prediction").as[Double].collect()
resultsUsingRaw2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
@@ -422,7 +422,7 @@ class LogisticRegressionSuite
}
// force it to use probability2prediction
- model.setRawPredictionCol("")
+ model.setRawPredictionCol("").setProbabilityCol("probability")
val resultsUsingProb2Predict =
model.transform(smallBinaryDataset).select("prediction").as[Double].collect()
resultsUsingProb2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
--
GitLab