Use withColumn in transform

feynmanliang · Jul 1, 2015 · 93aa2ff · 93aa2ff
1 parent d3a79f7
commit 93aa2ff
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
@@ -115,10 +115,7 @@ class KMeansModel private[ml] (
   }
 
   override def transform(dataset: DataFrame): DataFrame = {
-    dataset.select(
-          dataset("*"),
-          callUDF(predict _, IntegerType, col($(featuresCol))).as($(predictionCol))
-        )
+    dataset.withColumn($(predictionCol), callUDF(predict _, IntegerType, col($(featuresCol))))
   }
 
   override def transformSchema(schema: StructType): StructType = {