如何在 Pyspark 中聚集成对的 ID?

How to cluster pairs of IDs in Pyspark?

我有一列显示不同的 ID 对:

|column1|column2|
|123|456|
|123|678|
|678|315|
|333|777|

我正在尝试根据 ID 之间的重叠,通过添加家庭 ID 来对这些 ID 进行聚类:

|column1|column2|column3
|123|456|1|
|123|678|1|
|678|315|1|
|333|777|2|

关于如何做到这一点有什么建议吗?

我调查了以下问题,这对我描述此案例非常有用: see this question

现在我正在考虑在 Pyspark 中修复此问题的方法。这里有什么建议吗?

提前致谢!

正如您的链接问题所说,您需要一些图聚类算法。幸运的是,有一个 graphframes 包可以为您完成一切。要在 pyspark-shell 上使用 graphframes,您只需执行

pyspark --packages graphframes:graphframes:0.8.1-spark3.0-s_2.12

以下是如何解决您的问题的示例:

df.show()
+-------+-------+
|column1|column2|
+-------+-------+
|    123|    456|
|    123|    678|
|    678|    315|
|    333|    777|
+-------+-------+

from graphframes import GraphFrame
vertices = df.select('column1').union(df.select('column2')).distinct().toDF('id')
edges = df.toDF('src', 'dst')
g = GraphFrame(vertices, edges)

sc.setCheckpointDir('/tmp')
cc = g.connectedComponents().cache()

results = df.join(cc, df.column1 == cc.id).drop('id')
results.show()
+-------+-------+---------+
|column1|column2|component|
+-------+-------+---------+
|    333|    777|      333|
|    678|    315|      123|
|    123|    678|      123|
|    123|    456|      123|
+-------+-------+---------+