pyspark中的熊猫cumcount

Question

当前正在尝试将我从熊猫制作的脚本转换为pyspark，我有一个数据框，其中包含以下形式的数据：

index | letter
------|-------
0     | a
1     | a
2     | b
3     | c
4     | a
5     | a
6     | b

我想创建以下数据帧，其中存储每个字母实例的出现次数，例如，我们第一次看到“ a”时出现次数为0，第二次为1，第三次为2：

index | letter | occurrence
------|--------|-----------
0     | a      | 0
1     | a      | 1
2     | b      | 0
3     | c      | 0
4     | a      | 2
5     | a      | 3
6     | b      | 1

我可以使用以下方法在熊猫中实现这一目标

df['occurrence'] = df.groupby('letter').cumcount()

我将如何在pyspark中进行此操作？找不到类似的现有方法。

Answer 1

您要查找的功能称为window functions

from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

df.withColumn("occurence", row_number().over(Window.partitionBy("letter").orderBy("index")))