python - 为 pyspark 中的唯一行生成序列列

Question

我有一个 pyspark 数据框，如下所示:

col1 | col2 | col3
 R      a      abc
 R      a      abc
 G      b      def
 G      b      def
 G      b      def

我想向其中添加新列，它将根据如下计数为这些唯一行生成索引:

col1 | col2 | col3 | new_column
 R      a      abc      0
 R      a      abc      1
 G      b      def      0
 G      b      def      1
 G      b      def      2

请帮助我使用 pyspark 生成这个新专栏。

提前致谢!

Answer 1

按列对数据进行分区，然后使用行号为新列赋值。

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

windowSpec = W.partitionBy("col1","col2","col3").orderBy("col1","col2","col3")
df.withColumn("new_column", F.row_number().over(windowSpec)).show()