Include these Spark Window Functions in your Data Science Workflow | by Rahul Agarwal | Towards Data Science

In my last few posts on Spark, I explained how to work with PySpark RDDs and Dataframes.
Although these posts explain a lot on how to work with RDDs and Dataframe operations, and I would ask readers to go through them if they want to learn Spark Basics, I still didn’t mention quite a lot when it comes to working with PySpark Dataframes.
One such thing is the Spark window functions.
Recently, when I was working on one of my projects did I realize their power, and so this post is going to be about some of the most important Window functions available in Spark.
I will be working with the Data Science for COVID-19 in South Korea, which is one of the most detailed datasets on the internet for COVID.
Please note that I will be using this dataset to showcase the window functions, but this should not be in any way considered a data exploration exercise for this fantastic dataset.
For this post, I will use the TimeProvince dataframe, which contains daily case information for each province.
You can get rank as well as dense_rank on a group using this function. For example, you may want to have a column in your timeprovince table that provides the rank of each day in a province. We can do this by:
from pyspark.sql.window import Windowfrom pyspark.sql import functions as FwindowSpec = Window().partitionBy(['province']).orderBy(F.desc('confirmed'))timeprovince.withColumn("rank",F.rank().over(windowSpec)).show()
Here, we start by creating a window which is partitioned by province and ordered by the descending count of confirmed cases. One can begin to think of a window as a group of rows for a particular province in the order provided by the user. We can then add the rank easily by using the Rank function over this window, as shown above.
Sometimes our data science models may need lag based features. For example, a model might have variables like the price_last_week or sales_ quantity_previous_day. We can create such features using the lag function with window functions. Here I am trying to get the confirmed cases 7 days before. I am filtering to show the results as the first few days of corona cases were zeros. You can see here that the lag_7 day feature is shifted by seven days.
from pyspark.sql.window import WindowwindowSpec = Window().partitionBy(['province']).orderBy('date')timeprovinceWithLag = timeprovince.withColumn("lag_7",F.lag("confirmed", 7).over(windowSpec))timeprovinceWithLag.filter(timeprovinceWithLag.date>'2020-03-10').show()
Sometimes it helps to provide rolling averages to our models. For example, we might want to have a rolling 7-day sales sum/mean as a feature for our sales regression model. Let us calculate the rolling mean of confirmed cases for the last seven days here. This is what a lot of the people are already doing with this dataset to see the real trends.
from pyspark.sql.window import WindowwindowSpec = Window().partitionBy(['province']).orderBy('date').rowsBetween(-6,0)timeprovinceWithRoll = timeprovince.withColumn("roll_7_confirmed",F.mean("confirmed").over(windowSpec))timeprovinceWithRoll.filter(timeprovinceWithLag.date>'2020-03-10').show()
There are a few things here to understand. First is the rowsBetween(-6,0) function that we are using here. This function has a form of rowsBetween(start,end) with both start and end inclusive. Using this, we only look at the past seven days in a particular window, including the current_day. Here 0 specifies the current_row, and -6 specifies the seventh row previous to current_row. Remember, we count starting from 0.
So to get roll_7_confirmed for date 2020–03–22 we look at the confirmed cases for dates 2020–03–22 to 2020–03–16 and take their mean.
If we had used rowsBetween(-7,-1) we would just have looked at the past seven days of data and not the current_day.
One could also find a use for rowsBetween(Window.unboundedPreceding, Window.currentRow) where we take the rows between the first row in a window and the current_row to get running totals. I am calculating cumulative_confirmed here.
from pyspark.sql.window import WindowwindowSpec = Window().partitionBy(['province']).orderBy('date').rowsBetween(Window.unboundedPreceding,Window.currentRow)timeprovinceWithRoll = timeprovince.withColumn("cumulative_confirmed",F.sum("confirmed").over(windowSpec))timeprovinceWithRoll.filter(timeprovinceWithLag.date>'2020-03-10').show()
Here I have tried to summarize a few of the window functions I use while working with Spark. There are a lot of other ways you might think of using these window functions, but I hope this post gave you some clarity on how to work with them without being too confusing.
You can find all the code for this post at the GitHub repository.
Also, if you need to learn Spark basics, take a look at my previous post:
towardsdatascience.com
Also, if you want to learn more about Spark and Spark DataFrames, I would like to call out these excellent courses on Big Data Essentials: HDFS, MapReduce, and Spark RDD and Big Data Analysis: Hive, Spark SQL, DataFrames and GraphFrames by Yandex on Coursera.
Thanks for the read. I am going to be writing more beginner-friendly posts in the future too. Follow me up at Medium or Subscribe to my blog to be informed about them. As always, I welcome feedback and constructive criticism and can be reached on Twitter @mlwhiz
Also, a small disclaimer — There might be some affiliate links in this post to relevant resources, as sharing knowledge is never a bad idea.