How to get sequential ids in pyspark

Viewed 306

I have a pyspark dataframe with ids that repeat and are nonsequential. I would like to add a column of sequential id's, i.e. the second column below

id | seq
1  |1
3  |2
7  |3
3  |2
3  |2

The only way I have found to accomplish this is

window = Window.orderBy(col('id'))
df1 = df.select('id').distinct().withColumn('seq', F.row_number().over(window))
df.join(df1, on='id')

But that seems like not the best way. Is there another, quicker way to accomplish this?

1 Answers

Use dense_rank window function.

Example:

from pyspark.sql.window import *
from pyspark.sql.functions import *
w=Window.orderBy('id')

df.show()
#+---+
#| id|
#+---+
#|  1|
#|  3|
#|  3|
#|  3|
#|  7|
#+---+
df.withColumn("seq",dense_rank().over(w)).show()
#+---+---+
#| id|seq|
#+---+---+
#|  1|  1|
#|  3|  2|
#|  3|  2|
#|  3|  2|
#|  7|  3|
#+---+---+
Related