关联标识符对分组算法实现：基于PySpark/SQL处理1.8亿+行数据

我有一个包含"a"、"b"两列的数据库，每一行的两个标识符存在关联关系，目标是找到合适的算法，为每组互相关联的标识符分配统一的组编号。

示例输入表 #

期望输出结果 #

分组逻辑说明：1与2关联，2与4关联，4与5关联，5与1关联，因此1、2、4、5属于同一分组。

请问有没有适合在Spark（优先PySpark）或SQL中实现的算法？数据库中共有1.8亿+行数据，要求执行时间尽可能短。我已经尝试过Python实现的算法，但执行时间过长无法满足需求。

现有实现说明 #

目前使用的单机Python代码如下：

# Dict with result, for each id in 'a' -> a group id
result_dict = {}
group = 1

# Loop over the number of rows of the table
for i in range(len(data[0])):
    if data[0][i] in result_dict:
        # Important condition to change the group id off all ids already linked in previous iteration
        # In our example, to handle the couple (4,5) that appears before the (2,4)
        if (data[1][i] in result_dict) and (result_dict.get(data[0][i]) != result_dict.get(data[1][i])):
            temp = result_dict.get(data[1][i])
            for a, b in result_dict.items():
                if b == temp:
                    result_dict[a] = result_dict.get(data[0][i])
        else :
            result_dict[data[1][i]] = result_dict.get(data[0][i])
    elif data[0][i] not in result_dict and data[1][i] in result_dict:
        result_dict[data[0][i]] = result_dict.get(data[1][i])
    else :
        result_dict[data[0][i]] = group
        result_dict[data[1][i]] = group
        group = group + 1
        
# Fill the group column in the array        
for i in range(len(data[0])) : 
    data[2][i] = result_dict.get(data[1][i])

该代码存在大量循环逻辑，面对存储在HDFS上的1.8亿+行数据时，执行时间过长无法满足要求。

解决方案 #

这是典型的无向图连通分量计算问题，每一行的a和b对应图的一条无向边，你需要的Group就是每个连通分量的统一编号，以下是适配不同场景的实现方案：

PySpark实现（推荐，适配亿级数据） #

直接使用Spark生态的GraphFrames库自带的分布式连通分量算法，已经做了深度优化，性能远高于单机循环实现：

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

# 配置优化项，根据集群资源调整
spark.conf.set("spark.sql.shuffle.partitions", 200)
spark.sparkContext.setCheckpointDir("/tmp/checkpoint") # 避免迭代栈溢出

# 1. 读取原始数据，df为包含a、b两列的DataFrame
df = spark.read.parquet("你的HDFS数据路径")

# 2. 构造顶点表：所有a、b列的唯一值作为图的顶点
vertices = df.select(F.col("a").alias("id")).union(df.select(F.col("b").alias("id"))).distinct()

# 3. 构造边表：a为边起点，b为边终点
edges = df.select(F.col("a").alias("src"), F.col("b").alias("dst"))

# 4. 计算连通分量
g = GraphFrame(vertices, edges)
component_df = g.connectedComponents()

# 5. 关联回原始表得到最终结果
final_df = df.join(component_df, df.a == component_df.id, "left")\
             .select("a", "b", F.col("component").alias("Group"))

# 6. 结果写入HDFS
final_df.write.parquet("输出路径")

1.8亿行数据在8台节点的标准Spark集群上，通常可以在30分钟内完成计算。如果无法安装GraphFrames，也可以自行实现迭代式标签传播逻辑，每轮把节点的Group更新为最小的邻居Group，直到所有Group不再变化即可停止迭代。

SQL实现（仅适配小数据量，不推荐1.8亿行场景） #

如果只能使用SQL，可以用递归CTE实现，性能低于Spark分布式实现，仅作为备选：

WITH RECURSIVE all_edges AS (
  -- 把无向边转成双向边方便遍历
  SELECT a, b FROM your_table
  UNION ALL
  SELECT b, a FROM your_table
),
component AS (
  -- 初始每个节点的Group为自身的最小值
  SELECT a AS id, min(a) AS grp FROM all_edges GROUP BY a
  UNION ALL
  -- 迭代更新Group为更小的邻居Group
  SELECT e.a, c.grp
  FROM all_edges e
  JOIN component c ON e.b = c.id
  WHERE c.grp < (SELECT grp FROM component WHERE id = e.a)
)
-- 关联回原始表得到最终结果
SELECT t.a, t.b, min(c.grp) AS Group
FROM your_table t
JOIN (SELECT id, grp FROM component GROUP BY id, grp) c ON t.a = c.id
GROUP BY t.a, t.b

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