Создать столбец последовательности для уникальных строк в 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()