如何基于account_id-date-cluster_id映射表构建account_id-user_id映射DataFrame（PySpark）

我有一张映射表（account_id - date - cluster_id），示例如下：

规则说明 #

• 同一cluster_id在指定日期代表同一用户，但单个account_id在不同日期可对应多个cluster_id
• 指定日期的一个cluster_id可对应多个account_id
• 同一cluster_id可出现在多个日期，但不代表是同一用户

期望输出 #

需要基于此表创建account_id - user_id映射表，结果如下：

注：实际user_id的具体数值不重要，只要正确关联对应账户即可。

尝试的代码 #

我写了一段代码，但不确定是否正确，希望得到可行的解决方案：

df_user_mapping = cluster_df.withColumn('intermediate_group_id', F.dense_rank().over(Window.orderBy('year','month','cluster_id')).selectExpr('account_id','month','year','cluster_id as group_id','intermediate_group_id')
inter_groups = df_user_mapping.groupBy('account_id').agg(F.collect_set('intermediate_group_id').alias('intermediate_groups')))
groups = df_user_mapping.groupBy('account_id','year','month').agg(F.collect_set('group_id').alias('groups')))
joined = groups.join(inter_groups, ['account_id'], 'outer')
inter_explode = joined.select('account_id','groups','year','month',F.explode('intermediate_groups').alias('intermediate_group_id'))
group_explode = inter_explode.select('account_id','intermediate_group_id','year','month',F.explode('groups').alias('group_id'))
first_ids = group_explode.withColumn('first_id_inter', F.min('account_id').over(Window.partitionBy('inter_group_id'))).withColumn('first_id_group', F.min('first_id_inter').over(Window.partitionBy('group_id','year','month'))).withColumn('first_id', F.min('first_id_group').over(Window.partitionBy('account_id')))
final = first_ids.selectExpr('account_id','first_id as user_id').distinct()

初始模板代码 #

example = [('1',1,2018,'A'),
           ('2',1,2018,'A'),
           ('2',1,2019,'B'),
           ('3',1,2019,'B'),
           ('1',1,2020,'C'),
           ('4',1,2020,'C'),
           ('1',1,2021,'D'),
           ('5',1,2021,'A'),
          ]
example_df = spark.createDataFrame(example, ['account_id','month','year','cluster_id'])

这个问题本质是连通分量问题：通过cluster_id在不同日期的关联，把所有属于同一用户的account_id归为一组。比如account_id 1和2在2018年同属cluster A，2和3在2019年同属cluster B，1和4在2020年同属cluster C，因此1、2、3、4属于同一个连通分量；而account_id 5在2021年的cluster A和其他组无关联，单独成组。

方案一：使用GraphFrames（高效推荐） #

Spark的GraphFrames库可以快速处理连通分量问题，适合大数据场景。

1. 环境准备

提交Spark作业时需添加GraphFrames依赖：

spark-submit --packages graphframes:graphframes:0.8.2-spark3.2-s_2.12 your_script.py

2. 完整代码实现

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from graphframes import GraphFrame

# 初始化Spark会话
spark = SparkSession.builder.appName("AccountToUserIdMapping").getOrCreate()

# 加载示例数据
example = [('1',1,2018,'A'),
           ('2',1,2018,'A'),
           ('2',1,2019,'B'),
           ('3',1,2019,'B'),
           ('1',1,2020,'C'),
           ('4',1,2020,'C'),
           ('1',1,2021,'D'),
           ('5',1,2021,'A'),
          ]
example_df = spark.createDataFrame(example, ['account_id','month','year','cluster_id'])

# 步骤1：生成边表——同一日期同一cluster_id下的account_id两两相连
grouped = example_df.groupBy(F.concat_ws("_", "year", "month", "cluster_id").alias("group_key")) \
                   .agg(F.collect_list("account_id").alias("accounts"))

edges = grouped.withColumn("src", F.explode("accounts")) \
               .withColumn("dst", F.explode("accounts")) \
               .filter(F.col("src") != F.col("dst")) \
               .select("src", "dst")

# 步骤2：生成顶点表——所有唯一的account_id
vertices = example_df.select("account_id").distinct().withColumnRenamed("account_id", "id")

# 步骤3：构建图并计算连通分量
g = GraphFrame(vertices, edges)
connected_components = g.connectedComponents()

# 步骤4：将连通分量ID映射为简洁的user_id
final_mapping = connected_components.groupBy("component") \
                                   .agg(F.min("id").alias("user_id")) \
                                   .join(connected_components, on="component") \
                                   .select("id", "user_id") \
                                   .withColumnRenamed("id", "account_id")

# 展示结果
final_mapping.orderBy("account_id").show()

代码说明

1. 边表生成：同一(date, cluster_id)组内的account_id两两建立边，确保关联的账号被纳入同一个连通分量。
2. 顶点表生成：提取所有唯一account_id作为图的顶点。
3. 连通分量计算：connectedComponents方法自动识别所有连通的顶点组，每个组对应一个用户。
4. 映射简化：用组内最小的account_id作为user_id，满足关联要求且格式简洁。

输出结果

+----------+-------+
|account_id|user_id|
+----------+-------+
|         1|      1|
|         2|      1|
|         3|      1|
|         4|      1|
|         5|      5|
+----------+-------+

方案二：递归CTE（无需额外依赖） #

如果无法使用GraphFrames，可通过Spark SQL的递归CTE实现连通分量计算，适合小数据量场景：

from pyspark.sql import SparkSession

# 初始化Spark会话
spark = SparkSession.builder.appName("AccountToUserIdMapping").getOrCreate()

# 加载示例数据
example = [('1',1,2018,'A'),
           ('2',1,2018,'A'),
           ('2',1,2019,'B'),
           ('3',1,2019,'B'),
           ('1',1,2020,'C'),
           ('4',1,2020,'C'),
           ('1',1,2021,'D'),
           ('5',1,2021,'A'),
          ]
example_df = spark.createDataFrame(example, ['account_id','month','year','cluster_id'])
example_df.createOrReplaceTempView("example_df")

# 生成account_id与临时组的关联表
example_df.groupBy(F.concat_ws("_", "year", "month", "cluster_id").alias("temp_group")) \
          .agg(F.collect_list("account_id").alias("accounts")) \
          .select(F.explode("accounts").alias("account_id"), "temp_group") \
          .createOrReplaceTempView("account_group")

# 递归CTE计算连通分量
spark.sql("""
WITH RECURSIVE account_connections AS (
    SELECT account_id, account_id as root_id
    FROM (SELECT DISTINCT account_id FROM example_df) t
    UNION ALL
    SELECT ac.account_id, MIN(ac2.root_id) as root_id
    FROM account_connections ac
    JOIN account_group ag ON ac.account_id = ag.account_id
    JOIN account_group ag2 ON ag.temp_group = ag2.temp_group
    JOIN account_connections ac2 ON ag2.account_id = ac2.account_id
    WHERE ac.root_id > ac2.root_id
    GROUP BY ac.account_id
)
SELECT DISTINCT account_id, root_id as user_id
FROM account_connections
ORDER BY account_id
""").show()

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