根據列表中的匹配值在 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|
+-------+-----+