如何聚合Spark数据框以使用Scala获取稀疏向量?

时间:2017-07-24 16:10:25

标签: scala apache-spark spark-dataframe

我在Spark中有一个类似下面的数据框,我希望按id列对其进行分组,然后对于分组数据中的每一行,我需要创建一个包含来自{的元素的稀疏向量{1}}列指定的索引处的{1}}列。稀疏向量的长度是已知的,例如,对于此示例,为1000。

数据框weight

index

我读过this,这与我想做的有点相似,但对于一个rdd。有没有人知道使用Scala在Spark中为数据框执行此操作的好方法?

到目前为止,我的尝试是首先将权重和指数收集为如下列表:

df

看起来像:

+-----+------+-----+
|   id|weight|index|
+-----+------+-----+
|11830|     1|    8|
|11113|     1|    3|
| 1081|     1|    3|
| 2654|     1|    3|
|10633|     1|    3|
|11830|     1|   28|
|11351|     1|   12|
| 2737|     1|   26|
|11113|     3|    2|
| 6590|     1|    2|
+-----+------+-----+

然后我定义了一个udf并做了类似的事情:

val dfWithLists = df
    .groupBy("id")
    .agg(collect_list("weight") as "weights", collect_list("index") as "indices"))

但这似乎不起作用,感觉应该有一种更简单的方法来做到这一点,而不需要首先收集权重和索引到列表。

我对Spark,Dataframes和Scala都很陌生,因此非常感谢任何帮助。

1 个答案:

答案 0 :(得分:6)

你必须收集它们,因为向量必须是本地的,单机:https://spark.apache.org/docs/latest/mllib-data-types.html#local-vector

为了创建稀疏向量,您有2个选项,使用无序(索引,值)对或指定索引和值数组: https://spark.apache.org/docs/latest/api/scala/index.html#org.apache.spark.mllib.linalg.Vectors$

如果您可以将数据转换为其他格式(透视),您还可以使用VectorAssembler: https://spark.apache.org/docs/latest/ml-features.html#vectorassembler

通过一些小的调整,您可以使您的方法有效:

:paste
// Entering paste mode (ctrl-D to finish)

import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.mllib.regression.LabeledPoint

val df = Seq((11830,1,8), (11113, 1, 3), (1081, 1,3), (2654, 1, 3), (10633, 1, 3), (11830, 1, 28), (11351, 1, 12), (2737, 1, 26), (11113, 3, 2), (6590, 1, 2)).toDF("id", "weight", "index")

val dfWithFeat = df
  .rdd
  .map(r => (r.getInt(0), (r.getInt(2), r.getInt(1).toDouble)))
  .groupByKey()
  .map(r => LabeledPoint(r._1, Vectors.sparse(1000, r._2.toSeq)))
  .toDS

dfWithFeat.printSchema
dfWithFeat.show(10, false)


// Exiting paste mode, now interpreting.

root
|-- label: double (nullable = true)
|-- features: vector (nullable = true)

+-------+-----------------------+
|label  |features               |
+-------+-----------------------+
|11113.0|(1000,[2,3],[3.0,1.0]) |
|2737.0 |(1000,[26],[1.0])      |
|10633.0|(1000,[3],[1.0])       |
|1081.0 |(1000,[3],[1.0])       |
|6590.0 |(1000,[2],[1.0])       |
|11830.0|(1000,[8,28],[1.0,1.0])|
|2654.0 |(1000,[3],[1.0])       |
|11351.0|(1000,[12],[1.0])      |
+-------+-----------------------+

dfWithFeat: org.apache.spark.sql.Dataset[org.apache.spark.mllib.regression.LabeledPoint] = [label: double, features: vector]