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

Question

我有一个如下所示的 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 使用 phone 和 address 从上面的 spark 数据框中的连接组件。所以这个数据框可以看作是图的vertices数据框。

我想知道创建图表的 edges 数据框以馈送到graphframes 中的connectedComponents() 函数的最佳方法是什么？

理想情况下，edges 数据框应如下所示：

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

最后，connectedComponents() 结果应该如下所示。 id 0 & 1 基于 same_phone 关系在同一个组件中，1 & 3 基于 1 & 3 strong>same_address 关系。然后，这会将 4 作为另一个组件，与其他 ids 没有任何联系。

+---+-------------------+
|id |connected_component|
+---+-------------------+
|0  |1                  |
|1  |2                  |
|2  |1                  |
|3  |2                  |
|4  |3                  |
+---+-------------------+

提前致谢！

Answer 1

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|
+---+-------------------+