如何使用 Pyspark 中的 Graphframes 和 Spark Dataframe 中的原始数据获取连接组件?

How to Get Connected Component with Graphframes in Pyspark and Raw Data in Spark Dataframe?

我有一个 spark 数据框,如下所示:

+--+-----+---------+
|id|phone|  address|
+--+-----+---------+
| 0|  123| james st|
| 1|  177|avenue st|
| 2|  123|spring st|
| 3|  999|avenue st|
| 4|  678|  5th ave|
+--+-----+---------+

我正在尝试使用 graphframes 包来识别 ids 的连接组件,使用 phoneaddress 来自上面的 spark 数据框。所以这个数据框可以被视为图的个顶点个数据框。

我想知道创建图的 edges 数据框以输入 graphframes 中的 connectedComponents() 函数的最佳方法是什么?

理想情况下,edges 数据框应如下所示:

+---+---+------------+
|src|dst|relationship|
+---+---+------------+
| 0 |  2|  same_phone|
| 1 |  3|same_address|
+---+---+------------+

最后,connectedComponents() 结果应该如下所示。 id 0 & 1 基于 same_phone 关系和 1 & 3 基于 same_address 关系。然后,这将使 4 成为另一个与其他 ids.

没有联系的组件
+---+-------------------+
|id |connected_component|
+---+-------------------+
|0  |1                  |
|1  |2                  |
|2  |1                  |
|3  |2                  |
|4  |3                  |
+---+-------------------+

提前致谢!

from functools import reduce

edges = reduce(
    lambda x, y: x.union(y),
    [df.alias('t1')
       .join(df.alias('t2'), c)
       .filter('t1.id < t2.id')
       .selectExpr('t1.id src', 't2.id dst', "'same_%s' relationship"% c) for c in df.columns[1:]
    ]
)

edges.show()
+---+---+------------+
|src|dst|relationship|
+---+---+------------+
|  0|  2|  same_phone|
|  1|  3|same_address|
+---+---+------------+
import pyspark.sql.functions as F
from pyspark.sql.window import Window

connect = edges.select(
    F.array_sort(F.array('src', 'dst')).alias('arr')
).distinct().union(
    df.join(edges, (df.id == edges.src) | (df.id == edges.dst), 'anti').select(F.array('id'))
).withColumn(
    'connected_component', 
    F.row_number().over(Window.orderBy('arr'))
).select(F.explode('arr').alias('id'), 'connected_component')

connect.show()
+---+-------------------+
| id|connected_component|
+---+-------------------+
|  0|                  1|
|  2|                  1|
|  1|                  2|
|  3|                  2|
|  4|                  3|
+---+-------------------+