将任意数量的列转换为向量答案

【问题标题】：Convert arbitrary number of columns to Vector将任意数量的列转换为向量
【发布时间】：2018-03-31 20:19:29
【问题描述】：

如何将一组任意列转换为 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)

【问题讨论】：

标签： scala apache-spark apache-spark-sql dataset apache-spark-mllib

【解决方案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]|
+----+-------------+

【讨论】：

【解决方案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]|
// +----+-------------+

【讨论】：