How to create an index column, every 7 occurrences in pyspark - apache-spark

Question

I have a pyspark dataframe that looks like this:
import pandas as pd
foo = pd.DataFrame({'date_col':['2010-02-27','2010-01-20','2010-01-20','2010-01-21','2010-01-21','2010-02-21','2010-02-22','2010-02-23','2010-02-24','2010-02-25','2010-02-26','2010-01-20','2010-01-21','2010-02-20'], 'group':['a','a','a','a','a','a','a','a','a','a','a','b','b','b']})
I would like to create a week column, which would be an index, which increases every 7 ordered unique values of the date_col by group.
The resulting dataframe should look like this:
foo = pd.DataFrame({'date_col':['2010-02-27','2010-01-20','2010-01-20','2010-01-21','2010-01-21','2010-02-21','2010-02-22','2010-02-23','2010-02-24','2010-02-25','2010-02-26','2010-01-20','2010-01-21','2010-02-20'], 'group':['a','a','a','a','a','a','a','a','a','a','a','b','b','b'],
'week':[2,2,1,1,1,1,1,1,1,1,1,1,1,1]})
Any ideas could I do that in pyspark?
UPDATE
Some more explanation on the logic.
Basically the operation could be split into the following steps:
Order the foo on date_col grouped by group
Create a temp_index column, which would rank the date_col by group
Create a week column which would be the div of temp_index with 7

Answer 1

You can use dense_rank and divide the rank by 7. You need to subtract 1 before dividing because in SQL, ranks start from 1 rather than 0.
from pyspark.sql import functions as F, Window
df2 = df.withColumn(
'week',
(
(F.dense_rank().over(Window.partitionBy('group').orderBy('date_col')) - 1) / 7
).cast('int')
)
df2.show()
+----------+-----+----+
| date_col|group|week|
+----------+-----+----+
|2010-01-20| b| 0|
|2010-01-21| b| 0|
|2010-02-20| b| 0|
|2010-01-20| a| 0|
|2010-01-20| a| 0|
|2010-01-21| a| 0|
|2010-01-21| a| 0|
|2010-02-21| a| 0|
|2010-02-22| a| 0|
|2010-02-23| a| 0|
|2010-02-24| a| 0|
|2010-02-25| a| 0|
|2010-02-26| a| 1|
|2010-02-27| a| 1|
+----------+-----+----+