[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.
This commit is contained in:
parent
cca945b6aa
commit
dfc4c935ba
|
@ -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 {
|
||||
|
|
Loading…
Reference in a new issue