Skip to content
Snippets Groups Projects
Commit dfc4c935 authored by Yanbo Liang's avatar Yanbo Liang
Browse files

[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.
parent cca945b6
No related branches found
No related tags found
No related merge requests found
...@@ -360,7 +360,7 @@ class LogisticRegressionSuite ...@@ -360,7 +360,7 @@ class LogisticRegressionSuite
} }
// force it to use raw2prediction // force it to use raw2prediction
model.setProbabilityCol("") model.setRawPredictionCol("rawPrediction").setProbabilityCol("")
val resultsUsingRaw2Predict = val resultsUsingRaw2Predict =
model.transform(smallMultinomialDataset).select("prediction").as[Double].collect() model.transform(smallMultinomialDataset).select("prediction").as[Double].collect()
resultsUsingRaw2Predict.zip(results.select("prediction").as[Double].collect()).foreach { resultsUsingRaw2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
...@@ -368,7 +368,7 @@ class LogisticRegressionSuite ...@@ -368,7 +368,7 @@ class LogisticRegressionSuite
} }
// force it to use probability2prediction // force it to use probability2prediction
model.setRawPredictionCol("") model.setRawPredictionCol("").setProbabilityCol("probability")
val resultsUsingProb2Predict = val resultsUsingProb2Predict =
model.transform(smallMultinomialDataset).select("prediction").as[Double].collect() model.transform(smallMultinomialDataset).select("prediction").as[Double].collect()
resultsUsingProb2Predict.zip(results.select("prediction").as[Double].collect()).foreach { resultsUsingProb2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
...@@ -414,7 +414,7 @@ class LogisticRegressionSuite ...@@ -414,7 +414,7 @@ class LogisticRegressionSuite
} }
// force it to use raw2prediction // force it to use raw2prediction
model.setProbabilityCol("") model.setRawPredictionCol("rawPrediction").setProbabilityCol("")
val resultsUsingRaw2Predict = val resultsUsingRaw2Predict =
model.transform(smallBinaryDataset).select("prediction").as[Double].collect() model.transform(smallBinaryDataset).select("prediction").as[Double].collect()
resultsUsingRaw2Predict.zip(results.select("prediction").as[Double].collect()).foreach { resultsUsingRaw2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
...@@ -422,7 +422,7 @@ class LogisticRegressionSuite ...@@ -422,7 +422,7 @@ class LogisticRegressionSuite
} }
// force it to use probability2prediction // force it to use probability2prediction
model.setRawPredictionCol("") model.setRawPredictionCol("").setProbabilityCol("probability")
val resultsUsingProb2Predict = val resultsUsingProb2Predict =
model.transform(smallBinaryDataset).select("prediction").as[Double].collect() model.transform(smallBinaryDataset).select("prediction").as[Double].collect()
resultsUsingProb2Predict.zip(results.select("prediction").as[Double].collect()).foreach { resultsUsingProb2Predict.zip(results.select("prediction").as[Double].collect()).foreach {
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment