为 Spark 上的 fpgrowth 准备二进制表示的数据

Question

我目前正在使用 Kaggle 的 Santander 产品推荐数据集对 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()