PySpark - Window Functions

Introduction

Window functions in PySpark are a powerful feature for data manipulation and analysis. They allow you to perform complex calculations on subsets of data within a DataFrame, without the need for expensive joins or subqueries. In this tutorial, we will show you how to use window functions in PySpark. In order to do this, we will use the window functions row_number(), rank() and dense_rank() of PySpark with a certain window specification.

Window Functions

At its core, a window function performs a calculation on a set of rows. This "window" of rows is defined by a partition and an ordering within that partition. PySpark provides a rich set of window functions that can be applied to DataFrame columns.

Import Libraries

First, we import the following python modules:

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

Create SparkSession

Before we can work with Pyspark, we need to create a SparkSession. A SparkSession is the entry point into all functionalities of Spark.

In order to create a basic SparkSession programmatically, we use the following command:

spark = SparkSession \
    .builder \
    .appName("Python PySpark Example") \
    .getOrCreate()

Create PySpark DataFrame

Next, we create the PySpark DataFrame with some example data from a list. To do this, we use the method createDataFrame() and pass the data and the column names as arguments.

column_names = ["language", "framework", "users"]
data = [
    ("Python", "Django", 20000),
    ("Python", "FastAPI", 20000),
    ("JavaScript", "AngularJS", 5000),
    ("JavaScript", "ReactJS", 7000),
    ("Python", "Flask", 9000)
]
df = spark.createDataFrame(data, column_names)
df.show()

Define Window

Before applying window functions, it's essential to define a window specification. The window specification specifies how rows are partitioned and ordered within each partition.

In our example, we want to partition the data by the column "language" and order each partition in descending order based on the column "users". We can do this with the following command:

window_spec = Window.partitionBy("language").orderBy(F.desc("users"))