[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:
Yanbo Liang 2017-01-05 18:59:49 -08:00
parent cca945b6aa
commit dfc4c935ba

View file

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