Question

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

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，它有 1 行，并显示与他们共享标签的所有其他用户的数组：

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

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

2 个答案: