Spark SQL 102 — Aggregations and Window Functions | by David Vrba | Towards Data Science

Data aggregation is an important step in many data analyses. It is a way how to reduce the dataset and compute various metrics, statistics, and other characteristics. A related but slightly more advanced topic are window functions that allow computing also other analytical and ranking functions on the data based on a window with a so-called frame.
This is a continuation of a recent article in which we described what is a DataFrame and how transformations work in Spark SQL in general. Here we will dive into aggregations and window functions which are two specific groups of transformations, they are closely related, but as we will see, there is an important difference between them which is good to understand.
For the code, we will use PySpark API, the current version at the time of writing which is 3.1.2 (June 2021).
Let's start with the most simple aggregations which are computations in which we reduce the entire dataset to a single number. This might be like the total count of rows in the DataFrame or the sum/average of values in some specific column. For this purpose, we can use agg() function directly on the DataFrame and pass the aggregation functions as arguments in a comma-separated way:
from pyspark.sql.functions import count, sumdf.agg(count('*'))df.agg(count('*'), sum('price'))df.agg(  count('*').alias('number_of_rows'),  sum('price').alias('total_price'))
Notice that the output of the first example is a DataFrame with a single row and single column — it is just a number represented by a DataFrame. In the second example, the output is a DataFrame with a single row and two columns — one column for each aggregation function. In the last example, we can see that each of the aggregations can be also renamed using the alias() function.
Very often, we need to compute the aggregation, not for the entire DataFrame, but separately for each group of rows where the group is defined as rows that have the same value in a specific column. As an example, imagine a dataset of credit card transactions where each row is a unique transaction but different rows can belong to the same user (cardholder). Here, it might be useful to compute the aggregation separately for each user, and for this kind of aggregation we can use the groupBy transformation:
(  df  .groupBy('user_id')  .agg(count('*').alias('number_of_transactions')))
Again, we are using here the agg function and we can pass in any aggregation function such as count, countDistinct, sum, avg/mean, min, max, first, last, collect_list, collect_set, approx_count_distinct, corr, and for the complete list, I recommend to check the documentation. Notice that if we don’t rename the result of the aggregation, it will have a default name, which in the case of the count function is count(1).
Alternatively, we can call an aggregation function directly after the groupBy as follows
(  df  .groupBy('user_id')  .count())
Notice that using this syntax has the disadvantage that you can not rename the result of the aggregation directly using the alias() because here the count function returns back a DataFrame, so the alias would be applied on the entire DataFrame. So the renaming would have to be handled by another transformation such as withColumnRenamed(‘count’, ‘new_name’). Also, here you can call only one aggregation function at a time while the syntax with agg allows you to call as many functions as you want at the same time. For other functions that can be called after groupBy see the documentation.
One important property of these groupBy transformations is that the output DataFrame will contain only the columns that were specified as arguments in the groupBy() and the results of the aggregation. So if we call df.groupBy(‘user_id’).count(), no matter how many fields the df has, the output will have only two columns, namely user_id and count. Also, the number of rows represented by the output DataFrame will be smaller, or in a marginal case, the same as in the original df. The marginal case corresponds to a situation in which the grouping column has all values distinct so each group has exactly one row. As we will see, this will be different with window functions.
The window functions are a group of functions that can be called also over a group of rows similarly as we have seen in the previous case. There is a couple of differences, however. Firstly, after calling a window function, the dataset will not be reduced — all rows and all columns will be in the output DataFrame and the calculation will be added in a new column. The group of rows on which the function will be applied is again given by a specific column (or a list of columns) for which the rows have the same value and the group is referred to as a window. Also, the window functions are more flexible in the sense that sometimes you don’t want to apply the function on the entire window, but rather only on a subset of rows from the window — a so-called frame. Finally, the window can be also sorted because some functions (so-called ranking functions) require it. Let’s see the syntax for the window functions:
from pyspark.sql import Windoww = Window().partitionBy('user_id')df.withColumn('number_of_transactions', count('*').over(w))
As you can see, we first define the window using the function partitonBy() — this is analogous to the groupBy(), all rows that will have the same value in the specified column (here user_id) will form one window. Then we add a new column to the DataFrame and here we call our window function over the specified window.
This is a specific group of window functions that require the window to be sorted. As a specific example, consider the function row_number() that tells you the number of the row within the window:
from pyspark.sql.functions import row_numberw = Window.partitionBy('user_id').orderBy('transaction_date')df.withColumn('r', row_number().over(w))
Other ranking functions are for example rank() or dense_rank().
As mentioned above some functions can be applied over a subset of rows from the window. A typical use-case is to compute a cumulative sum of values where the frame would specify that we want to apply the function from the beginning of the window until the current row. Also, it is obvious that the order of the rows in the frame will be important because the cumulative sum will have a different shape if the rows go in a different order. The frame can be specified using one of these two functions:
rowsBetween()
rangeBetween()
Both of these two functions take two arguments: start and end of the frame and they can be specified as follows:
Window.unboundedPreceding, Window.unboundedFollowing — the entire window from the beginning to the end
Window.unboundedPreceding, Window.currentRow — from the beginning of the window to the current row, this is used for the cumulative sum
using numerical values, for example, 0 means currentRow, but the meaning of other values can differ based on the framing function rowsBetween/rangeBetween.
To understand the difference between rowsBetween and rangeBetween, let’s see the following example in which we have three columns, user_id, activity, and day and we want to sum the activity for each user:
df.withColumn('activity_sum', sum('activity').over(w))
We will sort the window by day and use the interval (-1, 0) to specify the frame. We will see that the interval has a different meaning for both functions. On the right side of the image you can see what will be the result of the sum in each case (for the sake of simplicity, we display in the image only one window for the user with id 100):
In the case of rowsBetween, on each row, we sum the activities from the current row and the previous one (if it exists), that’s what the interval (-1, 0) means. On the other hand, in the case of rangeBetween, on each row, we first need to compute the range of rows that will be summed by subtracting the value 1 from the value in the day column. For example on the 3rd row, we have 7–1 = 6 leading to the interval (6, 7) and we should sum up all rows where the value in the day column fits into this interval and in our case, it is only the current row, because there is no row with day=6.
As you can see, in the case of rangeBetween, the column by which we sort needs to be of some numerical type because Spark needs to do some arithmetic on the values in this column. This restriction is not present in the case of rowsBetween.
To mention some other window functions, see for example:
lead()
lag()
ntile()
nth_value()
cume_dist()
sorting the window will change the frame — this might not be intuitive — computing sum over a window that is sorted will lead to a different result as compared to a window that is not sorted. For more details, check my other article where I show an example.
Both transformations, groupBy, and window will require a specific partitioning, and if the partitioning is not present they will induce a shuffle — the data needs to be reorganized in such a way that all rows within one group/window need to be placed within one partition. The difference however is that with groupBy, Spark will partly aggregate the data first and then shuffle the reduced dataset as compared to window where the entire dataset will be shuffled.
If we don’t pass any argument to the partitionBy function and specify the window as w = Window().partitionBy(), the entire dataset will become one big window and all data will be shuffled to a single partition which may lead to performance issues because all data will be placed on a single node in the cluster.
In this article, we covered aggregation and window functions which are very frequently used transformations, especially among data analysts. We discussed what is the difference between calling groupBy and Window.partitionBy and we have seen different options for how to specify the frame on a window.