Generate sequence column for unique rows in pyspark

Viewed 2338

I have a pyspark dataframe which looks like below:

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

and I want to add new column to it which will generate the index for these unique rows as per the count like below:

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

Please help me to generate this new column using pyspark.

Thanks in advance!

1 Answers

Partition the data by column and then use the row number to assign values to new column.

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()
Related