时间戳

首页 > 解决方案 > 如何从 Pyspark 中的 spark 数据框创建边缘列表？

python - 如何从 Pyspark 中的 spark 数据框创建边缘列表？

问题描述

我graphframes在 pyspark 中使用某些图形类型的分析，想知道从顶点数据框创建边列表数据框的最佳方法是什么。

例如，下面是我的顶点数据框。我有一个 id 列表，它们属于不同的组。

+---+-----+
|id |group|
+---+-----+
|a  |1    |
|b  |2    |
|c  |1    |
|d  |2    |
|e  |3    |
|a  |3    |
|f  |1    |
+---+-----+

我的目标是创建一个边缘列表数据框来指示出现在常见组中的 id。请注意，1 个 id 可能出现在多个组中（例如，上面的 id a 在第 1 组和第 3 组中）。下面是我想获得的边缘列表数据框：

+---+-----+-----+
|src|dst  |group|
+---+-----+-----+
|a  |c    |1    |
|a  |f    |1    |
|c  |f    |1    |
|b  |d    |2    |
|a  |e    |3    |
+---+-----+-----+

提前致谢！

标签： pythonapache-sparkpysparkapache-spark-sqlgraphframes

解决方案

编辑 1

不确定这是否是更好的解决方法，但我做了一个解决方法：

import pyspark.sql.functions as f

df = df.withColumn('match', f.collect_set('id').over(Window.partitionBy('group')))

df = df.select(f.col('id').alias('src'),
               f.explode('match').alias('dst'),
               f.col('group'))

df = df.withColumn('duplicate_edges', f.array_sort(f.array('src', 'dst')))
df = (df
      .where(f.col('src') != f.col('dst'))
      .drop_duplicates(subset=['duplicate_edges'])
      .drop('duplicate_edges'))

df.sort('group', 'src', 'dst').show()

输出

+---+---+-----+
|src|dst|group|
+---+---+-----+
|  a|  c|    1|
|  a|  f|    1|
|  c|  f|    1|
|  b|  d|    2|
|  e|  a|    3|
+---+---+-----+

原始答案

尝试这个：

import pyspark.sql.functions as f

df = (df
      .groupby('group')
      .agg(f.first('id').alias('src'),
           f.last('id').alias('dst')))

df.show()

输出：

+-----+---+---+
|group|src|dst|
+-----+---+---+
|    1|  a|  c|
|    3|  e|  a|
|    2|  b|  d|
+-----+---+---+

推荐阅读