根据列表中的匹配值在 pyspark 数据框中添加新列

Question

我正在寻找有关 pyspark 的帮助，以添加具有匹配列表值的新列。

我有一个带有变量unique_ids的值列表

[Row(card_id=1), Row(card_id=2)]

对于列表中的每个值，如果列表值与列值匹配，则计算与该值匹配的行数并添加然后创建一个具有计数值的新列

这就是我获取列表的方式

unique_ids = data.select('card_id').distinct().collect()

示例 df

卡号
1
1
2
1
2
1

所需的数据框

卡号	数数
1	4
1	4
2	2
1	4
2	2
1	4

谢谢

Answer 1

使用窗函数count

import pyspark.sql.functions as F
from pyspark.sql.window import Window

unique_ids = data.withColumn('count', F.count('card_id').over(Window.partitionBy('card_id')))
unique_ids.show()

+-------+-----+
|card_id|count|
+-------+-----+
|      1|    4|
|      1|    4|
|      1|    4|
|      1|    4|
|      2|    2|
|      2|    2|
+-------+-----+