{"id":737,"date":"2020-08-30T15:25:58","date_gmt":"2020-08-30T15:25:58","guid":{"rendered":"https:\/\/justinmatters.co.uk\/wp\/?p=737"},"modified":"2020-09-01T15:46:49","modified_gmt":"2020-09-01T15:46:49","slug":"pyspark-window-functions","status":"publish","type":"post","link":"https:\/\/justinmatters.co.uk\/wp\/pyspark-window-functions\/","title":{"rendered":"PySpark Window Functions"},"content":{"rendered":"

PySpark window functions are useful when you want to examine relationships within groups of data rather than between groups of data as for groupBy. To use them you start by defining a window function, then select a separate function or set of functions to operate within that window.<\/p>\n

\"\"<\/p>\n

If you prefer to work through the example, you can download the workbook from Github<\/a>. The workbook is designed to work on Databricks Community Edition<\/a>. Alternatively for the next six months this link should work to give you direct access to a copy of the workbook on Databricks<\/a> since happily their public publishing system works now.<\/p>\n

Lets take a look at the sorts of things that can be achieved with window functions of varying complexity. The imports and dataframe\u00a0 I am going to be using for this explanation are:<\/p>\n

\r\nimport pandas as pd\r\nimport pyspark.sql.functions as fn\r\nfrom pyspark.sql import SparkSession\r\nfrom pyspark.sql import Window\r\n\r\n# Create a spark session\r\nspark_session = SparkSession.builder.getOrCreate()\r\n\r\n# lets define a demonstration DataFrame to work on\r\ndf_data = {\r\n  'partition': ['a','a', 'a', 'a', 'b', 'b', 'b', 'c', 'c',],\r\n  'col_1': [1,1,1,1,2,2,2,3,3,], \r\n  'aggregation': [1,2,3,4,5,6,7,8,9,],\r\n  'ranking': [4,3,2,1,1,1,3,1,5,],\r\n  'lagging': [9,8,7,6,5,4,3,2,1,],\r\n  'cumulative': [1,2,4,6,1,1,1,20,30,],\r\n}\r\ndf_pandas = pd.DataFrame.from_dict(df_data)\r\n# create spark dataframe\r\ndf = spark_session.createDataFrame(df_pandas)\r\n<\/pre>\n
Simple aggregation functions<\/h2>\n
we can use the standard group by aggregations with window functions. These functions use the simplest form of window which just defines grouping. However instead of producing one value for each group, we get the value added to every row of each partition in the window.<\/p>\n
\r\n# aggregation functions use the simplest form of window \r\n# which just defines grouping\r\naggregation_window = Window.partitionBy('partition')\r\n# then we can use this window function for our aggregations\r\ndf_aggregations = df.select(\r\n  'partition', 'aggregation'\r\n).withColumn(\r\n  'aggregation_sum', fn.sum('aggregation').over(aggregation_window),\r\n).withColumn(\r\n  'aggregation_avg', fn.avg('aggregation').over(aggregation_window),\r\n).withColumn(\r\n  'aggregation_min', fn.min('aggregation').over(aggregation_window),\r\n).withColumn(\r\n  'aggregation_max', fn.max('aggregation').over(aggregation_window),\r\n)\r\n<\/pre>\n
Row wise ordering and ranking functions<\/h2>\n
We can also use window funtions to order and rank data. These functions add an element to the definition of the window which defines both grouping and ordering<\/p>\n
\r\n# lets define a ranking window\r\nranking_window = Window.partitionBy('partition').orderBy('ranking')\r\n\r\ndf_ranks = df.select(\r\n  'partition', 'ranking'\r\n).withColumn(\r\n# note that fn.row_number() does not take any arguments\r\n  'ranking_row_number', fn.row_number().over(ranking_window)\r\n).withColumn(\r\n  # rank will leave spaces in ranking to account for preceding rows\r\n  # receiving equal ranks\r\n  'ranking_rank', fn.rank().over(ranking_window)\r\n).withColumn(\r\n  # dense rank does not account for previous equal rankings\r\n  'ranking_dense_rank', fn.dense_rank().over(ranking_window)\r\n).withColumn(\r\n  # percent rank ranges between 0-1 not 0-100\r\n  'ranking_percent_rank', fn.percent_rank().over(ranking_window)\r\n).withColumn(\r\n  # fn.ntile takes a parameter for now many 'buckets' to divide \r\n  # rows into when ranking\r\n  'ranking_ntile_rank', fn.ntile(2).over(ranking_window)\r\n)\r\n<\/pre>\n
We can also revers the order of ranking using .desc(). In this case the window would be defined as follows<\/p>\n
\r\n# lets define a ranking window in reverse order using .desc() \r\n# note that we need to use fn.col to define the column \r\ndesc_ranking_window = Window.partitionBy(\r\n  'partition'\r\n).orderBy(\r\n  fn.col('ranking').desc()\r\n)\r\n<\/pre>\n
Cumulative Calculations (Running Totals and Averages)<\/h3>\n
There are often good reasons to want to create a running total or running average column. In some cases we might want running totals for subsets of data. Window functions can be useful for that sort of thing.<\/p>\n
In order to calculate such things we need to add yet another element to the window. Now we account for partition, order and which rows should be covered by the function. This can be done in two ways we can use rangeBetween<\/code> to define how similar values in the window must be to be considered, or we can use rowsBetween<\/code> to define how many rows should be considered. The current row is considered row zero, the following rows are numbered positively and the preceding rows negatively. For cumulative calculations you can define “all previous rows” with Window.unboundedPreceding<\/code> and “all following rows” with Window.unboundedFolowing<\/code><\/p>\n
Note that the window may vary in size as it progresses over the rows since at the start and end part of the window may “extend past” the existing rows<\/p>\n
\r\n#suppose we want to average over the previous, current and next values\r\n# running calculations need a more complicated window as shown here\r\ncumulative_window_1 = Window.partitionBy(\r\n  'partition'\r\n).orderBy(\r\n  'cumulative'\r\n  # for a rolling average lets use rowsBetween\r\n).rowsBetween(\r\n  -1,1\r\n)\r\n\r\ndf_cumulative_1 = df.select(\r\n  'partition', 'cumulative'\r\n).withColumn(\r\n  'cumulative_avg', fn.avg('cumulative').over(cumulative_window_1)\r\n)\r\n\r\n# running totals also require a more complicated window as here.\r\ncumulative_window_2 = Window.partitionBy(\r\n  'partition'\r\n).orderBy(\r\n  'cumulative'\r\n  # in this case we will use rangeBetween for the sum\r\n).rangeBetween(\r\n  # here we use Window.unboundedPreceding to catch all earlier rows\r\n  Window.unboundedPreceding, 0\r\n)\r\n\r\ndf_cumulative_2 = df.select(\r\n  'partition', 'cumulative'\r\n).withColumn(\r\n  'cumulative_sum', fn.sum('cumulative').over(cumulative_window_2)\r\n)\r\n<\/pre>\n
Combining Windows and Calling Different Columns<\/h2>\n
It is also possible to combine windows and also to call windows on columns other than the ordering column. These more advanced uses can require careful thought to ensure you achieve the intended results.<\/p>\n
First lets look at using multiple window functions in a single expression<\/p>\n
\r\n# we can make a window function equivalent to a standard groupBy:\r\n\r\n# first define two windows\r\naggregation_window = Window.partitionBy('partition')\r\ngrouping_window = Window.partitionBy('partition').orderBy('aggregation')\r\n\r\n# then we can use this window function for our aggregations\r\ndf_aggregations = df.select(\r\n  'partition', 'aggregation'\r\n).withColumn(\r\n  # note that we calculate row number over the grouping_window\r\n  'group_rank', fn.row_number().over(grouping_window) \r\n).withColumn(\r\n  # but we calculate other columns over the aggregation_window\r\n  'aggregation_sum', fn.sum('aggregation').over(aggregation_window),\r\n).withColumn(\r\n  'aggregation_avg', fn.avg('aggregation').over(aggregation_window),\r\n).withColumn(\r\n  'aggregation_min', fn.min('aggregation').over(aggregation_window),\r\n).withColumn(\r\n  'aggregation_max', fn.max('aggregation').over(aggregation_window),\r\n).where(\r\n  fn.col('group_rank') == 1\r\n).select(\r\n  'partition', \r\n  'aggregation_sum', \r\n  'aggregation_avg', \r\n  'aggregation_min', \r\n  'aggregation_max'\r\n)\r\n\r\n# this is equivalent to the rather simpler expression below\r\ndf_groupby = df.select(\r\n  'partition', 'aggregation'\r\n).groupBy(\r\n  'partition'\r\n).agg(\r\n  fn.sum('aggregation').alias('aggregation_sum'),\r\n  fn.avg('aggregation').alias('aggregation_avg'),\r\n  fn.min('aggregation').alias('aggregation_min'),\r\n  fn.max('aggregation').alias('aggregation_max'),\r\n)\r\n<\/pre>\n
Secondly here is an example of ordering by one column but operating on a different column. This can only be done for calculation functions which take an column as a parameter<\/p>\n
\r\n# create a window on one column but use the window on another column \r\nlag_window = Window.partitionBy('partition').orderBy('lagging')\r\n\r\ndf_cumulative_2 = df.select(\r\n  'partition', 'lagging', 'cumulative',\r\n).withColumn(\r\n  'lag_the_laggging_col', fn.lag('lagging', 1).over(lag_window)\r\n).withColumn(\r\n  # It is possible to lag a column which was not the orderBy column\r\n  'lag_the_cumulative_col', fn.lag('cumulative', 1).over(lag_window)\r\n)\r\n<\/pre>\n
The effect of windows functions is best understood by experimenting with them so I encourage you to make use of the Databricks workbook linked at the top of the page<\/a>.<\/p>\n
Further Reading<\/p>\n
Further useful references about PySpark window functions can be found here:<\/p>\n