fix mllib

Author: cloud-fan

mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala

@@ -131,7 +131,7 @@ final class OneVsRestModel private[ml] (
 
     // output label and label metadata as prediction
     aggregatedDataset
-      .withColumn($(predictionCol), labelUDF(col(accColName)).as($(predictionCol), labelMetadata))
+      .withColumn($(predictionCol), labelUDF(col(accColName)), Some(labelMetadata))
       .drop(accColName)
   }
 
@@ -203,8 +203,8 @@ final class OneVsRest(override val uid: String)
       // TODO: use when ... otherwise after SPARK-7321 is merged
       val newLabelMeta = BinaryAttribute.defaultAttr.withName("label").toMetadata()
       val labelColName = "mc2b$" + index
-      val labelUDFWithNewMeta = labelUDF(col($(labelCol))).as(labelColName, newLabelMeta)
-      val trainingDataset = multiclassLabeled.withColumn(labelColName, labelUDFWithNewMeta)
+      val trainingDataset = multiclassLabeled
+        .withColumn(labelColName, labelUDF(col($(labelCol))), Some(newLabelMeta))
       val classifier = getClassifier
       val paramMap = new ParamMap()
       paramMap.put(classifier.labelCol -> labelColName)

mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala

@@ -75,7 +75,7 @@ final class Bucketizer(override val uid: String)
     }
     val newCol = bucketizer(dataset($(inputCol)))
     val newField = prepOutputField(dataset.schema)
-    dataset.withColumn($(outputCol), newCol.as($(outputCol), newField.metadata))
+    dataset.withColumn($(outputCol), newCol, Some(newField.metadata))
   }
 
   private def prepOutputField(schema: StructType): StructField = {

mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala

@@ -341,7 +341,7 @@ class VectorIndexerModel private[ml] (
     val newField = prepOutputField(dataset.schema)
     val transformUDF = udf { (vector: Vector) => transformFunc(vector) }
     val newCol = transformUDF(dataset($(inputCol)))
-    dataset.withColumn($(outputCol), newCol.as($(outputCol), newField.metadata))
+    dataset.withColumn($(outputCol), newCol, Some(newField.metadata))
   }
 
   override def transformSchema(schema: StructType): StructType = {

mllib/src/main/scala/org/apache/spark/ml/feature/VectorSlicer.scala

@@ -119,8 +119,7 @@ final class VectorSlicer(override val uid: String)
         case features: SparseVector => features.slice(inds)
       }
     }
-    dataset.withColumn($(outputCol),
-      slicer(dataset($(inputCol))).as($(outputCol), outputAttr.toMetadata()))
+    dataset.withColumn($(outputCol), slicer(dataset($(inputCol))), Some(outputAttr.toMetadata()))
   }
 
   /** Get the feature indices in order: indices, names */

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/AnalysisSuite.scala

@@ -133,5 +133,7 @@ class AnalysisSuite extends AnalysisTest {
     // CreateStruct is a special case that we should not trim Alias for it.
     plan = testRelation.select(CreateStruct(Seq(a, (a + 1).as("a+1"))).as("col"))
     checkAnalysis(plan, plan)
+    plan = testRelation.select(CreateStructUnsafe(Seq(a, (a + 1).as("a+1"))).as("col"))
+    checkAnalysis(plan, plan)
   }
 }

sql/core/src/main/scala/org/apache/spark/sql/DataFrame.scala

@@ -1135,17 +1135,18 @@ class DataFrame private[sql](
    * @group dfops
    * @since 1.3.0
    */
-  def withColumn(colName: String, col: Column): DataFrame = {
+  def withColumn(colName: String, col: Column, metadata: Option[Metadata] = None): DataFrame = {
     val resolver = sqlContext.analyzer.resolver
     val replaced = schema.exists(f => resolver(f.name, colName))
+    val aliasedColumn = metadata.map(md => col.as(colName, md)).getOrElse(col.as(colName))
     if (replaced) {
       val colNames = schema.map { field =>
         val name = field.name
-        if (resolver(name, colName)) col.as(colName) else Column(name)
+        if (resolver(name, colName)) aliasedColumn else Column(name)
       }
       select(colNames : _*)
     } else {
-      select(Column("*"), col.as(colName))
+      select(Column("*"), aliasedColumn)
     }
   }