Group into current row previous group values with Pyspark

user_6865633 0
Viewed 140

So, I have a pyspark dataframe organised in this way:

ID timestamp value1 value2
1 1 a x
2 1 a y
1 2 b x
2 2 b y
1 3 c y
2 3 d y
1 4 l y
2 4 s y

and let's say that the timestamp is the number of day from the beginning of time. What I'd like to do is, for each line, to group into a list the values up to -x days regarding the current ID, so to have:

ID timestamp value1 value2 list_value_1
1 1 a X a
2 1 a y a
1 2 b x a,b
2 2 b y a,b
1 3 c y a,b,c
2 3 d y a,b,d
1 3 c y b,c,l
2 3 d y b,d,s

I imagine I should do that with a Window but I'm not sure on how to proceed (I'm quite bad with Windows for some reason).

pyspark apache-spark-sql apache-spark python
1 Answers

You can do a collect_list over a Window betweeen the current row and two preceding rows, and combine the list into a comma-separated string using concat_ws:

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'list_value_1', 
    F.concat_ws(',', 
        F.collect_list('value1').over(
            Window.partitionBy('ID').orderBy('timestamp').rowsBetween(-2, 0)
        )
    )
)

df2.show()
+---+---------+------+------+------------+
| ID|timestamp|value1|value2|list_value_1|
+---+---------+------+------+------------+
|  1|        1|     a|     x|           a|
|  1|        2|     b|     x|         a,b|
|  1|        3|     c|     y|       a,b,c|
|  1|        4|     l|     y|       b,c,l|
|  2|        1|     a|     y|           a|
|  2|        2|     b|     y|         a,b|
|  2|        3|     d|     y|       a,b,d|
|  2|        4|     s|     y|       b,d,s|
+---+---------+------+------+------------+
user_141657300
Related