如何使用 Spark/Python 枚举组中的行？

Question

我想像 Pandas 一样枚举分组值：

枚举 DataFrame 中每个组的每一行

Spark/Python 中的方法是什么？

Answer 1

具有row_number窗口功能：

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

w = Window.partitionBy("some_column").orderBy("some_other_column")
df.withColumn("rn", row_number().over(w))