how to collect a list of new IDs between two successive dates with pyspark

user_10056100 0
Viewed 228

I am doing a groupby over a week column in a PYSPARK dataframe and collecting new different IDs between two sucessive weeks number.

I tried doing a groupby over week column then aggregating the collect_set method on the dataframe to get all the available IDs, then i compared the lists one by one to get the different ID between two successive lists by week number as below:

Input :

Week ID
1 ID_1
1 ID_2
1 ID_3
2 ID_1
2 ID_4
2 ID_5
3 ID_6

OUTPUT:

Week ID_List diff_list new_different_ID_count
1 ID_1, ID_2,ID_3 ID_1, ID_2, ID_3 -
2 ID_1, ID_4, ID_5 ID_4, ID_5 2
3 ID_2, ID_6 ID_6 1

The probleme here is that as the huge number of IDs i have in my case (more than 9 million ID) the spark session is being killed due to an Out of Memory Error i think (Error 500)!

Is there any other solution to get the list of new different IDs beetween two successive weeks with PYSPARK ?

pyspark apache-spark python
1 Answers

In order to scale, you need to aggregate by the ID instead, without collecting any results. Try the following approach:

import pyspark.sql.functions as F
from pyspark.sql import Window

data = spark.createDataFrame([(1, "ID_1"), (1, "ID_2"), (1, "ID_3"),
                              (2, "ID_1"), (2, "ID_4"), (2, "ID_5"),
                              (3, "ID_6")], ["Week", "ID"])

win = Window.partitionBy('ID').orderBy('Week')

agg_data = (
  data
  .withColumn("prevWeek", F.lag("Week", offset=1).over(win))
  .withColumn("isInPrevWeek", 
              F.col("prevWeek").isNotNull() & ((F.col("Week") - F.col("prevWeek")) == 1))
  .filter(~F.col("isInPrevWeek"))
  .groupBy("Week")
  .agg(F.count("*").alias("newIDs"),
       F.array_sort(F.collect_list("ID")).alias("showNewIDs"))  # Remove in production
  .orderBy("Week")
)

agg_data.show()

First, the function lag creates a new column with the previous Week using a Window function which allows to consider each ID individually (partition) and sorts the weeks chronologically. That scales well because the Spark tasks are made of group of IDs.

Then, isInPrevWeek checks if the ID was indeed in the previous week. If so, the record is filtered out. Now, you just need to count the remining IDs by week.

+----+------+------------------+
|Week|newIDs|        showNewIDs|
+----+------+------------------+
|   1|     3|[ID_1, ID_2, ID_3]|
|   2|     2|      [ID_4, ID_5]|
|   3|     1|            [ID_6]|
+----+------+------------------+

Note that the snippet collects the IDs just for illustrative purposes but that is not required for the computation.

user_3771180
Related