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