基于分组的另一列中的重叠条目查找一列中的公共对

Question

我需要查询方面的帮助。说我有一个像这样的数据框：

+------+------+
|userid|songid|
+------+------+
|     1|     a|
|     1|     b|
|     1|     c|
|     2|     a|
|     2|     d|
|     3|     c|
|     4|     e|
|     4|     d|
|     5|     b|
+------+------+

我想返回一个数据帧，该数据帧具有userid对，它们之间至少有一个共同的songid。上面的数据框看起来像这样：

+------+------+
|userid|friendid|
+------+------+
|     1|     2|
|     1|     3|
|     1|     5|
|     2|     4|
+------+------+

我该怎么做？

Answer 1

一种简单的方法是使用自我Join：

data = [(1, 'a'), (1, 'b'), (1, 'c'),
        (2, 'a'), (2, 'd'), (3, 'c'),
        (4, 'e'), (4, 'd'), (5, 'b')
        ]

df = spark.createDataFrame(data, ["userid", "songid"])

# join on songId = songId and userid different
join_condition = (col("u1.songid") == col("u2.songid")) & (col("u1.userid") != col("u2.userid"))

df.alias("u1").join(df.alias("u2"), join_condition, "inner") \
    .select(sort_array(array(col("u1.userid"), col("u2.userid"))).alias("pairs")) \
    .distinct() \
    .select(col("pairs").getItem(0).alias("userid"), col("pairs").getItem(1).alias("friendid"))\
    .show()

+------+--------+
|userid|friendid|
+------+--------+
|     1|       3|
|     1|       5|
|     2|       4|
|     1|       2|
+------+--------+