基于PySpark GraphFrame实现关联零件列值的连通分量分组

• 构造边数据时仅关联了colVal1列，colVal2~colVal5的关联关系完全没有纳入图计算，丢失了大量关联边
• 没有过滤空字符串的目标值，空值会作为无效节点参与计算，干扰连通分量结果
• 最后关联结果时使用了不存在的device字段作为关联键，关联逻辑完全错误

我们需要先把所有partnumber和对应5个属性列的非空关联都转换为边，再进行连通分量计算，具体代码如下：

from pyspark.sql.functions import lit, array, explode, col

# 第一步：构造全量有效边
# 把partnumber和每个colVal的非空关系都拆成src-dst边
edges = df_part_groups.select(
    col("partnumber").alias("src"),
    explode(array("colVal1", "colVal2", "colVal3", "colVal4", "colVal5")).alias("dst")
).filter(col("dst") != "") # 过滤空值的无效边

# 第二步：构造顶点，包含所有出现过的part编号
vertices = edges.select("src").distinct().union(edges.select("dst").distinct()).withColumnRenamed("src", "id")

# 第三步：计算连通分量
g = G.GraphFrame(vertices, edges)
# 连通分量计算需要先设置checkpoint目录，可根据实际情况调整路径
spark.sparkContext.setCheckpointDir("./tmp_checkpoint")
cc = g.connectedComponents()

# 第四步：关联回原表得到每个partnumber对应的分组ID
result = df_part_groups.join(
    cc,
    df_part_groups.partnumber == cc.id,
    "left"
).orderBy("component", "partnumber")

display(result)

计算完成后component列就是关联分组的唯一ID，相同分组的所有关联part会共享同一个component值，符合预期的分组输出。如果需要聚合每个分组下的所有part编号，可以按component分组后做collect_list聚合即可。

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