将任意数量的列转换为 Vector

Question

如何将一组任意列转换为 Mllib 向量？基本上，我有数据帧的第一列，具有固定名称，然后是多个任意命名的列，每个列内部都有双精度值。

这样：

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

这里有一种方法可以做到这一点：

import org.apache.spark.sql.functions._
import org.apache.spark.mllib.linalg.DenseVector
val ds = Seq(
  ("val1", 0.0, 1.0, 1.0),
  ("val2", 2.0, 1.0, 5.0)
).toDF("name", "a", "b", "c").
as[(String, Double, Double, Double)]
val colList = ds.columns
val keyCol = colList(0)
val valCols = colList.drop(1)
def arrToVec = udf(
  (s: Seq[Double]) => new DenseVector(s.toArray)
)
ds.select(
  col(keyCol), arrToVec( array(valCols.map(x => col(x)): _*) ).as("values")
).show
// +----+-------------+
// |name|       values|
// +----+-------------+
// |val1|[0.0,1.0,1.0]|
// |val2|[2.0,1.0,5.0]|
// +----+-------------+