将任意数量的列转换为Vector

Question

如何将一组任意列转换为Mllib Vector？ 基本上，我的DataFrame的第一列有一个固定的名称，然后是一些任意命名的列，每列都有Double值。

像这样：

name  |  a  |  b  |  c  |
val1  | 0.0 | 1.0 | 1.0 |
val2  | 2.0 | 1.0 | 5.0 |

可以是任意数量的列。我需要获得以下数据集：

final case class ValuesRow(name: String, values: Vector)

Answer 1

这可以使用VectorAssembler以简单的方式完成。要合并到Vector的列用作输入，在本例中是除第一列之外的所有列。

val df = spark.createDataFrame(Seq(("val1", 0, 1, 1), ("val2", 2, 1, 5)))
  .toDF("name", "a", "b", "c")

val columnNames = df.columns.drop(1) // drop the name column    
val assembler = new VectorAssembler()
  .setInputCols(columnNames)  
  .setOutputCol("values")

val df2 = assembler.transform(df).select("name", "values").as[ValuesRow]

结果将是包含名称和值列的数据集：

+----+-------------+
|name|       values|
+----+-------------+
|val1|[0.0,1.0,1.0]|
|val2|[2.0,1.0,5.0]|
+----+-------------+

Answer 2

这是一种方法：

rpart