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

Question

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

Enumerate each row for each group in a 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))

Answer 2

您可以通过执行以下操作在rdd级别实现此目的：

rdd = sc.parallelize(['a', 'b', 'c'])
df = spark.createDataFrame(rdd.zipWithIndex())
df.show()

结果如下： +---+---+ | _1| _2| +---+---+ | a| 0| | b| 1| | c| 2| +---+---+ 如果您只需要唯一ID，而不是真正的连续索引，您也可以使用 zipWithUniqueId()更高效，因为在每个分区上本地完成。