如何在PySpark中生成统计唯一ID数量的类别交叉表？

问题原因 #

你用df.groupby('Category').pivot('Category').count()得到错误结果，是因为这个逻辑是统计每个类别分组内各类型的出现行数，既没有去重同一个ID的重复类别，也没有统计同时属于两个类别的唯一ID数，完全不符合需求。

正确实现方法 #

针对1000万行的规模，推荐两种高效实现方式：

方法一：自连接+去重统计 #

这种方法逻辑直观，适合类别分布分散的场景：

1. 去重ID-Category组合：先去掉同一ID下重复的类别，避免后续重复计算

dedup_df = df.dropDuplicates(["ID", "Category"])

2. 自连接生成类别对：将去重后的表与自身按ID连接，得到每个ID对应的所有类别组合

cross_df = dedup_df.alias("a").join(dedup_df.alias("b"), on="ID", how="inner")

3. 分组透视生成交叉表：按行类别分组，列类别透视，统计唯一ID数量

from pyspark.sql.functions import col

cross_table = cross_df.groupBy(col("a.Category").alias("row_category")) \
    .pivot(col("b.Category")) \
    .agg(countDistinct("ID"))

方法二：收集类别集合+笛卡尔积展开 #

这种方法先聚合每个ID的类别集合，再生成组合，减少中间数据量，适合每个ID类别数量较少的场景：

1. 去重并收集每个ID的类别集合

from pyspark.sql.functions import collect_set

id_cats = df.dropDuplicates(["ID", "Category"]) \
    .groupBy("ID") \
    .agg(collect_set("Category").alias("categories"))

2. 生成类别笛卡尔积组合：将每个ID的类别集合与自身做笛卡尔积，得到所有类别对

from pyspark.sql.functions import explode

id_cross = id_cats.withColumn("col_category", explode(col("categories"))) \
    .withColumn("row_category", explode(col("categories")))

3. 透视生成交叉表

cross_table = id_cross.groupBy("row_category") \
    .pivot("col_category") \
    .agg(countDistinct("ID"))

示例验证 #

用你提供的示例数据测试：
原数据：

执行后得到的交叉表：

其中B行C列的1表示同时属于B和C的唯一ID（即ID2）数量，符合需求。

内容的提问来源于stack exchange，提问作者Mufeez