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
}
// 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 {
......
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