创建一个DataFrame，显示与每个ID共享一列值的其他ID

Question

我有以下DataFrame或user_id和label列。一个用户可以有多个标签。

df = spark.createDataFrame(
[(1, "a"), (2, "b"), (3, "a"), (1, "c"), (4, "b"), (5, "c"), (6, "a"), (7, "e")], ['user_id', 'label']
)
+-------+-----+                                                                 
|user_id|label|
+-------+-----+
|      1|    a|
|      2|    b|
|      3|    a|
|      1|    c|
|      4|    b|
|      5|    c|
|      6|    a|
|      7|    e|    
+-------+-----+

我想创建一个新的DataFrame，每个用户有一行，并显示他们共享标签的所有其他用户的数组：

+-------+-------------+
|user_id|  other_users|
+-------+-------------+
|      1|    [3, 5, 6]|
|      2|          [4]|
|      3|       [1, 6]|
|      4|          [2]|
|      5|          [1]|
|      6|       [1, 3]|
|      7|           []|
+-------+-------------+

实现这一目标的最佳方式是什么？

Answer 1

您可以加入数据帧本身并使用collect_list

from  pyspark.sql.functions import col, collect_list
df = (df
.join(df.selectExpr('user_id ui', 'label lb'),
[col('label') == col('lb'), col('user_id') != col('ui')],
'left')
.groupBy('user_id').agg(collect_list('ui').alias('other_users')))
df.show()
+-------+-----------+
|user_id|other_users|
+-------+-----------+
|      7|         []|
|      6|     [1, 3]|
|      5|        [1]|
|      1|  [5, 3, 6]|
|      3|     [1, 6]|
|      2|        [4]|
|      4|        [2]|
+-------+-----------+

Answer 2

另一种方式。我这样做了，但看到@Wai Ha Lee的回答，我投了赞成票，因为它更简洁。虽然有所保留，但已经决定以另一种方式分享和摆出姿势。

h=Window.partitionBy('label')#grouper 1
g=Window.partitionBy('user_id')#grouper 2
df1=(df.withColumn('other_users',F.collect_list('user_id').over(h))#For every lable collect user_id
.withColumn("user_id", array(df['user_id']))#Convert user_id column to list
.withColumn('other_users',F.array_distinct(F.flatten(F.collect_list('other_users').over(g))))#Combine user_id lists in the other_users columns
.withColumn("other_users", array_except(col("other_users"), col("user_id"))))#Exclude user_ids
df1.show()