將任意數量的列轉換爲矢量

Question

如何將一組任意列轉換爲Mllib矢量？ 基本上，我有我的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

下面是做這件事：

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]| 
// +----+-------------+