From b4badae8bce863558eaeb831293ac8bbfc87b179 Mon Sep 17 00:00:00 2001 From: Josh Howes Date: Sat, 18 Jun 2016 14:13:18 -0400 Subject: [PATCH] Rename output column to match predictionCol --- .../org/apache/spark/ml/classification/OneVsRest.scala | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala index 047a378b79aa..817ecb285a9a 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala @@ -196,8 +196,13 @@ final class OneVsRestModel private[ml] ( } // output label and label metadata as prediction + val predictionMetadata = new MetadataBuilder() + .withMetadata(labelMetadata) + .putString("name", predictionCol.name) + .build() + aggregatedDataset - .withColumn($(predictionCol), labelUDF(col(accColName)), labelMetadata) + .withColumn($(predictionCol), labelUDF(col(accColName)), predictionMetadata) .drop(accColName) }