如何用Spark DataFrame函数将二进制数据转换为FPGrowth所需格式？

当然可以用Spark DataFrame的内置函数轻松完成这个转换！下面提供两种实用的实现方法，你可以根据数据规模和场景选择：

方法一：使用高阶函数（推荐，无Shuffle更高效） #

这种方法直接利用Spark的when、array和filter高阶函数，不需要进行数据重组或Shuffle，性能更优：

from pyspark.sql import functions as F

# 假设你的原始DataFrame名为df
# 先提取所有商品列（排除id列）
item_columns = [col for col in df.columns if col != "id"]

# 转换生成items列
transformed_df = df.withColumn(
    "items",
    # 过滤掉数组中的null值，只保留值为1的列名
    F.filter(
        # 对每个商品列，值为1时返回列名，否则返回null
        F.array(*[F.when(F.col(col) == 1, F.lit(col)).otherwise(F.lit(None)) for col in item_columns]),
        lambda x: x.isNotNull()
    )
).select("id", "items")

# 查看结果
transformed_df.show()

代码解释： #

• 首先筛选出所有非id的列，也就是代表商品的二进制列；
• 对每个商品列，用when判断值是否为1：如果是，返回列名（用lit把列名转为字符串常量），否则返回null；
• 把这些结果组合成一个数组，再用filter剔除数组中的null值，剩下的就是该行用户拥有的商品列表；
• 最后只保留id和items列，得到FPGrowth需要的格式。

方法二：宽表转窄表再聚合（适合理解逻辑） #

如果更习惯宽表转长表的思路，也可以用explode和groupBy来实现：

from pyspark.sql import functions as F

item_columns = [col for col in df.columns if col != "id"]

# 第一步：把每个商品列转为(item, value)的struct，再展开成行
long_df = df.select(
    "id",
    F.explode(
        F.array(*[F.struct(F.lit(col).alias("item"), F.col(col).alias("value")) for col in item_columns])
    ).alias("temp")
)

# 第二步：过滤值为1的行，再按id聚合收集商品列表
transformed_df = long_df.filter(F.col("temp.value") == 1)\
                       .groupBy("id")\
                       .agg(F.collect_list("temp.item").alias("items"))

transformed_df.show()

代码解释： #

• 先把每个商品列包装成包含item（列名）和value（二进制值）的结构体，再用explode把数组展开成多行；
• 过滤出value=1的行，只保留用户拥有的商品记录；
• 最后按id分组，用collect_list把每个用户的商品收集成列表。

两种方法最终都会输出你需要的格式：

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

内容的提问来源于stack exchange，提问作者Alper Nebi Kanlı