apache-spark - 在 Spark 上为 fpgrowth 准备二进制表示的数据

Question

我目前正在研究来自 Kaggle 的 Santander Product Recommendation 数据集，以对 FPGrowth 进行实验。

来自 pyspark (ML) 的 FPGrowth 算法需要数据框作为项目集：

+---+------------+
| id|       items|
+---+------------+
|  0|   [A, B, E]|
|  1|[A, B, C, E]|
|  2|      [A, B]|
+---+------------+

但是我拥有的数据是这种格式：

+---+---+---+---+---+---+
| id|  A|  B|  C|  D|  E|
+---+---+---+---+---+---+
|  0|  1|  1|  0|  0|  1|
|  1|  1|  1|  1|  0|  1|
|  2|  1|  1|  0|  0|  0|
+---+---+---+---+---+---+

我试图通过用列名替换 1 并从中创建列表来解决它，但这不起作用。

有没有办法通过使用 Spark 数据框函数来执行这种转换？

非常感谢你！

Answer 1

使用 udf：

from pyspark.sql.functions import udf, struct

@udf("array<string>")
def as_basket(row):
    return [k for k, v in row.asDict().items() if v]

df.withColumn("basket", as_basket(struct(*df.columns[1:]))).show()