在 Spark 中使用 VectorAssembler

Question

得到如下dataframe（假定已经是dataframe）：

val df = sc.parallelize(Seq((1, 2, 10), (3, 4, 11), (5, 6, 12)))
           .toDF("a", "b", "c")

我想将列（不是全部）合并为一列，并使其成为 Array[Double] 的 rdd。我正在执行以下操作：

import org.apache.spark.ml.feature.VectorAssembler
val colSelected = List("a","b")
val assembler = new VectorAssembler()
    .setInputCols(colSelected.toArray)
    .setOutputCol("features")
val output = assembler.transform(df).select("features").rdd

到这里就OK了。现在输出是 RDD[spark.sql.Row] 格式的数据帧。我无法将其转换为 RDD[Array[Double]] 的格式。有什么办法吗？

我试过类似下面的方法但没有成功：

output.map { case Row(a: Vector[Double]) => a.getAs[Array[Double]]("features")}

Answer 1

正确的解决方案（假设 Spark 2.0+，在 1.x 中使用 o.a.s.mllib.linalg.Vector）：

import org.apache.spark.ml.linalg.Vector

output.map(_.getAs[Vector]("features").toArray)

ml / mllib Vector 由 VectorAssembler 创建与 scala.collection.Vector.
Row.getAs 应与预期类型一起使用。它不执行任何类型转换并且 o.a.s.ml(lib).linalg.Vector 不是 Array[Double].

在 Spark 中使用 VectorAssembler

Using VectorAssembler in Spark

apache-spark

apache-spark-ml