Github user zhengruifeng commented on a diff in the pull request:
https://github.com/apache/spark/pull/22991#discussion_r236110139
--- Diff:
mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala ---
@@ -219,14 +225,20 @@ final class OneVsRestModel private[ml] (
Vectors.dense(predArray)
}
- // output the index of the classifier with highest confidence as
prediction
- val labelUDF = udf { (rawPredictions: Vector) =>
rawPredictions.argmax.toDouble }
-
- // output confidence as raw prediction, label and label metadata as
prediction
- aggregatedDataset
- .withColumn(getRawPredictionCol, rawPredictionUDF(col(accColName)))
- .withColumn(getPredictionCol, labelUDF(col(getRawPredictionCol)),
labelMetadata)
- .drop(accColName)
+ if (getPredictionCol != "") {
--- End diff --
I implemented this in another way, classificationmodel update the output
dataset, and I direct return the output in each if clause.
Then I update the to follow ClassificationModel, and update the
outputColumns in each clauses. And `withColumns` is used to return the output
columns.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
