PySpark中关联ArrayType列与IntegerType列的分组需求

问题本质 #

你的需求核心是识别ID的连通分量：df2中的collected_grp数组定义了grp之间的关联关系，df1中的ID隶属于特定grp，因此需要先建立grp的连通关系，再将同一连通grp下的所有ID合并为一组。

完整实现代码 #

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

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

# 原始输入数据
data = [
    ['7933546917973caa8c2898c834446415', '3ef2e38d48a9af3e096ddd3bc3816afb', 1],
    ['7d693086c5b8f74cbe881166cf3c2a29', 'fcb907411aff4f44c599cf03d23327c0', 2],
    ['7e18b452bb1e2845800a71d9431033b6', '9bc9d06e0efb16abde20c35ba36a2f1b', 3],
    ['7e18b452bb1e2845800a71d9431033b6', 'ff351ada316cbb0f270f935adfd16ad4', 4],
    ['8240cf1e442a97aa91d1029270728bbb', '484f25e9ab91af2c116cd788c91bdc82', 5],
    ['8919d5fd5b6fd118c1c6b691c65c9df9', '8dc7dfb4466590375f1aaac7fc8cb987', 6],
    ['8919d5fd5b6fd118c1c6b691c65c9df9', '9b93e3cfc5605e74ce2ce4c9450fd622', 7],
    ['8dc7dfb4466590375f1aaac7fc8cb987', '9b93e3cfc5605e74ce2ce4c9450fd622', 8],
    ['8f459a7cff281bad73f604166841849e', '41f007c0cc45c228e246f1cc91145878', 9],
    ['99f70106443a6f3f5c69d99a49d22d01', 'be73ca52536d13dfea295d4fcd273fde', 10],
    ['a9781767ca4fe8fb1282ee003d2c06ac', 'cb6feb2f38731fc7832545cbe2ac881b', 11],
    ['f4901968c29e928fc7364411b03336d4', '6fa82a51f17f0bf258fe06befc661216', 12],
    ['f6da014449e6fa82c24d002b4a27b105', '41f007c0cc45c228e246f1cc91145878', 13],
    ['f6da014449e6fa82c24d002b4a27b105', '8f459a7cff281bad73f604166841849e', 14],
    ['f93c0028bb26bc9b99fca1db300c2ac1', 'ccce888c5813025e95434d7ceedf1db3', 15],
    ['ff351ada316cbb0f270f935adfd16ad4', '9bc9d06e0efb16abde20c35ba36a2f1b', 16],
    ['ffe20a2c61638bb10bf943c42b4d794f', '985e237162ccfc04874664648893c241', 17],
]

# 生成df1：包含所有ID与对应grp的映射
df1 = spark.createDataFrame(
    [(row[0], row[2]) for row in data] + [(row[1], row[2]) for row in data],
    ["ID", "grp"]
).dropDuplicates()

# 生成df2：按你提供的逻辑生成collected_grp
df = spark.createDataFrame(data, schema=['id1', 'id2', 'grp'])
df2 = df.alias('df1')\
    .join(df.alias('df2'), (F.col('df1.id1') == F.col('df2.id2')), 'left')\
    .select(F.array_distinct(F.array(F.col('df1.id1'), F.col('df1.id2'), F.col('df2.id1'), F.col('df2.id2'))).alias('ID'), F.col('df1.grp') )
df3 = df2.select(F.explode('ID').alias('ID'), 'grp').dropna()
df2 = df3.groupBy('ID').agg(F.collect_set('grp').alias('collected_grp'))

# 步骤1：提取grp的连通边
# 展开collected_grp生成所有grp关联对
edges_df = df2.select(F.explode(F.col("collected_grp")).alias("src"))\
    .crossJoin(df2.select(F.explode(F.col("collected_grp")).alias("dst")))\
    .filter(F.col("src") < F.col("dst"))\
    .dropDuplicates()
# 补充自环边，保证单个grp被识别为独立连通分量
vertices_df = df1.select(F.col("grp").alias("id")).dropDuplicates()
edges_df = edges_df.union(vertices_df.select(F.col("id").alias("src"), F.col("id").alias("dst")))

# 步骤2：用GraphFrames计算连通分量
g = GraphFrame(vertices_df, edges_df)
connected_components = g.connectedComponents()

# 步骤3：将ID映射到对应的连通组
id_component_map = df1.join(connected_components, df1.grp == connected_components.id, "left")\
    .select("ID", "component")

# 步骤4：按连通组分组合并ID
final_result = id_component_map.groupBy("component")\
    .agg(F.array_distinct(F.collect_list("ID")).alias("ID"))\
    .select("ID")

# 输出结果
final_result.show(truncate=False)

关键步骤说明 #

1. 构建完整ID-grp映射：将原始数据中的id1和id2都纳入df1，避免遗漏任何ID。
2. 提取grp连通关系：把df2中的数组展开为grp对，补充自环边确保单个grp的独立性。
3. 计算连通分量：借助GraphFrames的connectedComponents方法，高效识别所有连通的grp组。
4. 合并ID分组：将同一连通grp下的所有ID收集并去重，得到最终的ID分组。

替代方案（无GraphFrames时） #

如果无法使用GraphFrames，可通过迭代合并grp的方式实现，但仅适合小数据量：

# 简单迭代合并逻辑示例
grp_pairs = edges_df.collect()
parent = {}

def find(u):
    while parent[u] != u:
        parent[u] = parent[parent[u]]
        u = parent[u]
    return u

def union(u, v):
    u_root = find(u)
    v_root = find(v)
    if u_root != v_root:
        parent[v_root] = u_root

# 初始化父节点
for grp in vertices_df.collect():
    parent[grp.id] = grp.id

# 合并所有连通对
for pair in grp_pairs:
    union(pair.src, pair.dst)

# 生成grp到根的映射
grp_component = spark.createDataFrame([(k, find(k)) for k in parent.keys()], ["grp", "component"])
# 后续步骤同之前的id_component_map和final_result生成

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