Rename output column to match predictionCol

josh-howes · josh-howes · commit b4badae8bce8 · 2016-09-08T06:04:52.000-04:00
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
@@ -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)
   }
 

Original file line number	Diff line number	Diff line change
`@@ -196,8 +196,13 @@ final class OneVsRestModel private[ml] (`
`196`	`196`	`}`
`197`	`197`
`198`	`198`	`// output label and label metadata as prediction`
	`199`	`+ val predictionMetadata = new MetadataBuilder()`
	`200`	`+ .withMetadata(labelMetadata)`
	`201`	`+ .putString("name", predictionCol.name)`
	`202`	`+ .build()`
	`203`	`+`
`199`	`204`	`aggregatedDataset`
`200`		`- .withColumn($(predictionCol), labelUDF(col(accColName)), labelMetadata)`
	`205`	`+ .withColumn($(predictionCol), labelUDF(col(accColName)), predictionMetadata)`
`201`	`206`	`.drop(accColName)`
`202`	`207`	`}`
`203`	`208`