VectorToArray.scala

package org.apache.spark.ml.feature

import org.apache.spark.SparkException
import org.apache.spark.annotation.{Experimental, Since}
import org.apache.spark.ml.Transformer
import org.apache.spark.ml.attribute.{Attribute, NominalAttribute}
import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol}
import org.apache.spark.ml.param.{ParamMap, StringArrayParam}
import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable, Identifiable}
import org.apache.spark.mllib.linalg.{VectorUDT, SparseVector, DenseVector, Vector}
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._


class VectorToArray(override val uid: String)
  extends Transformer with HasInputCol with HasOutputCol with DefaultParamsWritable {

  def this() =
    this(Identifiable.randomUID("vectorToArray"))

  /** @group setParam */
  def setInputCol(value: String): this.type = set(inputCol, value)

  /** @group setParam */
  def setOutputCol(value: String): this.type = set(outputCol, value)

  override def transformSchema(schema: StructType): StructType = {
    val inputColName = $(inputCol)
    val inputDataType = schema(inputColName).dataType
    require(inputDataType.isInstanceOf[VectorUDT],
      s"The input column $inputColName must be a vector type, " +
        s"but got $inputDataType.")
    val inputFields = schema.fields
    val outputColName = $(outputCol)
    require(inputFields.forall(_.name != outputColName),
      s"Output column $outputColName already exists.")
    val outputFields = inputFields :+ StructField($(outputCol), ArrayType(DoubleType,false),true)
    StructType(outputFields)
  }

  override def transform(dataset: DataFrame): DataFrame = {

    val transformer = udf { vec: Vector  =>
        vec.toArray
    }

    val outputColName = $(outputCol)
    dataset.select(col("*"),
      transformer(dataset($(inputCol))).as(outputColName))
  }

  override def copy(extra: ParamMap): VectorToArray = {
    defaultCopy(extra)
  }
}

object VectorToArray extends DefaultParamsReadable[VectorToArray] {
  override def load(path: String): VectorToArray = super.load(path)
}