PySpark：从组中的值创建一个向量[重复]

Posted 2023-04-15

【中文标题】PySpark：从组中的值创建一个向量[重复]

【英文标题】：PySpark: create a vector from values in a group [duplicate]

【发布时间】：2019-08-02 03:15:12

【问题描述】：

我目前有一个用户交易历史数据集，格式如下：

+---------+------------+------------+
| user_id | order_date | product_id |
+---------+------------+------------+
|       1 |   20190101 |        123 |
|       1 |   20190102 |        331 |
|       1 |   20190301 |       1029 |
+---------+------------+------------+

我正在尝试转换要用于 Item2Vec 模型的数据集——我认为它必须看起来像这样：

+---------+-------------------+
| user_id |      seq_vec      |
+---------+-------------------+
|    1    |  [123, 331, 1029] |
-------------------------------

通过查看 Word2Vec (https://spark.apache.org/docs/2.2.0/ml-features.html#word2vec) 的示例，我假设数据集必须以这种方式格式化。

如果我按user_id 分组，是否有内置的 PySpark 方法可以根据product_id 列中的值创建向量？

【参考方案1】：

collect_list 成功了

import pyspark.sql.functions as F

rawData = [(1, 20190101, 123),
           (1, 20190102, 331),
           (1, 20190301, 1029)]

df = spark.createDataFrame(rawData).toDF("user_id", "order_date", "product_id")

df.groupBy("user_id").agg(F.collect_list("product_id").alias("vec")).show()

+-------+----------------+
|user_id|             vec|
+-------+----------------+
|      1|[123, 331, 1029]|
+-------+----------------+