Pyspark Windows Functions : Notes by Arun Nautiyal

There are three types of Windows functions:

In [4]: import findspark

findspark.init()

import pyspark

from pyspark.sql import SparkSession

In [6]: spark=SparkSession.builder.appName('Arun_test').getOrCreate()

In [34]: simpleData = [('Arun','UK',50000),

('Arjun','MP',30000),

('James','USA',20000),

('Kumar',"UK",35000),

('Ajay','UK',50000)]

columns = ['Name','State','Salary']

df = spark.createDataFrame(data=simpleData,schema=columns)

 In [35]: df.show(truncate=False)

+-----+-----+------+

|Name |State|Salary|

+-----+-----+------+

|Arun |UK |50000 |

|Arjun|MP |30000 |

|James|USA |20000 |

|Kumar|UK |35000 |

|Ajay |UK |50000 |

+-----+-----+------+

In [15]: df.printSchema()

root

|-- Name: string (nullable = true)

|-- State: string (nullable = true)

|-- Salary: long (nullable = true)

Ranking Functions
row_number()
rank()
dense_rank()
ntile()
percent_rank()

In [36]: from pyspark.sql.window import Window

from pyspark.sql.functions import row_number

from pyspark.sql.functions import col

In [37]: WindFunc = Window.partitionBy('State').orderBy(col('Salary').desc())

df.withColumn('RN',row_number().over(WindFunc)).show()

+-----+-----+------+---+

| Name|State|Salary| RN|

+-----+-----+------+---+

|James| USA| 20000| 1|

| Arun| UK| 50000| 1|

| Ajay| UK| 50000| 2|

|Kumar| UK| 35000| 3|

|Arjun| MP| 30000| 1|

+-----+-----+------+---+

In [38]: from pyspark.sql.functions import rank

df.withColumn('Person_Rank',rank().over(WindFunc)).show(truncate=False)

+-----+-----+------+-----------+

|Name |State|Salary|Person_Rank|

+-----+-----+------+-----------+

|James|USA |20000 |1 |

|Arun |UK |50000 |1 |

|Ajay |UK |50000 |1 |

|Kumar|UK |35000 |3 |

|Arjun|MP |30000 |1 |

+-----+-----+------+-----------+

In [39]: from pyspark.sql.functions import dense_rank

df.withColumn('Person_Dense_Rank',dense_rank().over(WindFunc)).show(truncate=False)

+-----+-----+------+-----------------+

|Name |State|Salary|Person_Dense_Rank|

+-----+-----+------+-----------------+

|James|USA |20000 |1 |

|Arun |UK |50000 |1 |

|Ajay |UK |50000 |1 |

|Kumar|UK |35000 |2 |

|Arjun|MP |30000 |1 |

+-----+-----+------+-----------------+

In [40]: from pyspark.sql.functions import percent_rank

df.withColumn('percent_rank',percent_rank().over(WindFunc)).show()

+-----+-----+------+------------+

| Name|State|Salary|percent_rank|

+-----+-----+------+------------+

|James| USA| 20000| 0.0|

| Arun| UK| 50000| 0.0|

| Ajay| UK| 50000| 0.0|

|Kumar| UK| 35000| 1.0|

|Arjun| MP| 30000| 0.0|

+-----+-----+------+------------+

In [43]: from pyspark.sql.functions import ntile

df.withColumn("ntile_values",ntile(3).over(WindFunc)).show()

+-----+-----+------+------------+

| Name|State|Salary|ntile_values|

+-----+-----+------+------------+

|James| USA| 20000| 1|

| Arun| UK| 50000| 1|

| Ajay| UK| 50000| 2|

|Kumar| UK| 35000| 3|

|Arjun| MP| 30000| 1|

+-----+-----+------+------------+

In [44]: df.withColumn("ntile_values",ntile(2).over(WindFunc)).show()

+-----+-----+------+------------+

| Name|State|Salary|ntile_values|

+-----+-----+------+------------+

|James| USA| 20000| 1|

| Arun| UK| 50000| 1|

| Ajay| UK| 50000| 1|

|Kumar| UK| 35000| 2|

|Arjun| MP| 30000| 1|

+-----+-----+------+------------+

 In [45]: df.withColumn("ntile_values",ntile(1).over(WindFunc)).show()

+-----+-----+------+------------+

| Name|State|Salary|ntile_values|

+-----+-----+------+------------+

|James| USA| 20000| 1|

| Arun| UK| 50000| 1|

| Ajay| UK| 50000| 1|

|Kumar| UK| 35000| 1|

|Arjun| MP| 30000| 1|

+-----+-----+------+------------+

PySpark Window Analytic functions:

In [47]: from pyspark.sql.functions import cume_dist

df.withColumn("cume_dist_value",cume_dist().over(WindFunc)).show(truncate= False)

+-----+-----+------+------------------+

|Name |State|Salary|cume_dist_value |

+-----+-----+------+------------------+

|James|USA |20000 |1.0 |

|Arun |UK |50000 |0.6666666666666666|

|Ajay |UK |50000 |0.6666666666666666|

|Kumar|UK |35000 |1.0 |

|Arjun|MP |30000 |1.0 |

+-----+-----+------+------------------+

In [48]: from pyspark.sql.functions import lag

df.withColumn("lag_value",lag('salary',1).over(WindFunc)).show()

+-----+-----+------+---------+

| Name|State|Salary|lag_value|

+-----+-----+------+---------+

|James| USA| 20000| null|

| Arun| UK| 50000| null|

| Ajay| UK| 50000| 50000|

|Kumar| UK| 35000| 50000|

|Arjun| MP| 30000| null|

+-----+-----+------+---------+

In [50]: # craeted a new window function without partition

WindFunc_new = Window.orderBy(col('Salary').desc())

df.withColumn('lag_val',lag('salary',1).over(WindFunc_new)).show()

+-----+-----+------+-------+

| Name|State|Salary|lag_val|

+-----+-----+------+-------+

| Arun| UK| 50000| null|

| Ajay| UK| 50000| 50000|

|Kumar| UK| 35000| 50000|

|Arjun| MP| 30000| 35000|

|James| USA| 20000| 30000|

+-----+-----+------+-------+

In [51]: from pyspark.sql.functions import lead

df.withColumn("load_value",lead("salary",1).over(WindFunc)).show(truncate=False)

+-----+-----+------+----------+

|Name |State|Salary|load_value|

+-----+-----+------+----------+

|James|USA |20000 |null |

|Arun |UK |50000 |50000 |

|Ajay |UK |50000 |35000 |

|Kumar|UK |35000 |null |

|Arjun|MP |30000 |null |

+-----+-----+------+----------+

In [53]: df.withColumn('leadValue_without_partition',lead('salary',1).over(WindFunc_new)).show(truncate=False)

df.withColumn('leadValue_without_partition',lead('salary',2).over(WindFunc_new)).show(truncate=False)

+-----+-----+------+---------------------------+

|Name |State|Salary|leadValue_without_partition|

+-----+-----+------+---------------------------+

|Arun |UK |50000 |50000 |

|Ajay |UK |50000 |35000 |

|Kumar|UK |35000 |30000 |

|Arjun|MP |30000 |20000 |

|James|USA |20000 |null |

+-----+-----+------+---------------------------+

+-----+-----+------+---------------------------+

|Name |State|Salary|leadValue_without_partition|

+-----+-----+------+---------------------------+

|Arun |UK |50000 |35000 |

|Ajay |UK |50000 |30000 |

|Kumar|UK |35000 |20000 |

|Arjun|MP |30000 |null |

|James|USA |20000 |null |

+-----+-----+------+---------------------------+

Pyspark Aggregate Functions

Note: In the below code, I have implemented the aggregate functions.

Please note

that I have used WindFunc(orderBy is mandatory) with row_number since it is a ranking function and windAgg(without orderBy) with the aggregate
functions

In [96]: windAgg = Window.partitionBy("state")

from pyspark.sql.functions import min, max, avg, sum, col, row_number

df.withColumn("row",row_number().over(WindFunc)) \

.withColumn("min",min(col('salary')).over(windAgg)) \

.withColumn('max',max(col('salary')).over(windAgg)) \

.withColumn('avg',avg('salary').over(windAgg)) \

.withColumn('sum',sum('salary').over(windAgg))\

.where(col('row') ==1) \

.select('state','min','max','avg','sum').show()

+-----+-----+-----+-------+------+

|state| min| max| avg| sum|

+-----+-----+-----+-------+------+

| USA|20000|20000|20000.0| 20000|

| UK|35000|50000|45000.0|135000|

| MP|30000|30000|30000.0| 30000|

+-----+-----+-----+-------+------+