DataFrame Window Functions

Performing advanced relational analytical calculations in PySpark across partitioned ranges using Window, partitionBy, orderBy, rank, lead, and lag.

What are Window Functions?

Window functions compute values over a partitioned range of rows called a window frame. Unlike a groupBy() aggregation which collapses rows to a single summary row, window functions compute a value for every single row while retaining the original rows in the DataFrame.

This is extremely powerful for running calculations like cumulative sums, patient vitals rankings, moving averages, and calculating differences between consecutive rows.

Syntax and Window Specifications

To use window functions, define a WindowSpec using the Window class, and then apply functions using the .over() method:

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

# Define the Window frame: partition by department, and order by salary descending
windowSpec = Window.partitionBy("department").orderBy(F.col("salary").desc())

# Apply ranking within each department
ranked_df = df.withColumn("rank", F.dense_rank().over(windowSpec))

Core Window Functions

• Ranking: row_number(), rank(), dense_rank()
• Analytics: lead(col, offset) (gets values from subsequent rows), lag(col, offset) (gets values from preceding rows)
• Aggregates over Window: sum(), avg(), min(), max()

Example Usage Pipeline

Below is a complete, copy-paste-ready PySpark script demonstrating window calculations:

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

# 1. Setup local Spark session
spark = SparkSession.builder \
    .appName("DataFrame Window Demo") \
    .master("local[*]") \
    .getOrCreate()

# 2. Dummy dataset (ICU patient heart rate telemetry readings over time)
data = [
    ("PAT_01", "10:00", 72),
    ("PAT_01", "10:05", 78),
    ("PAT_01", "10:10", 85),
    ("PAT_02", "10:00", 95),
    ("PAT_02", "10:05", 110),
    ("PAT_02", "10:10", 102),
]
columns = ["patient_id", "reading_time", "heart_rate"]
df = spark.createDataFrame(data, columns)

# 3. Create Window Specifications
# Window A: Partition by patient, ordered by reading time ascending
patient_time_window = Window.partitionBy("patient_id").orderBy("reading_time")

# 4. Apply Window Transformations:
# - Calculate the difference in heart rate between consecutive readings (current vs. previous) using lag()
# - Calculate the average heart rate of the patient over all readings in this window using avg()
analytical_df = df \
    .withColumn("prev_heart_rate", F.lag("heart_rate", 1).over(patient_time_window)) \
    .withColumn("heart_rate_change", F.col("heart_rate") - F.col("prev_heart_rate")) \
    .withColumn("patient_avg_hr", F.round(F.avg("heart_rate").over(Window.partitionBy("patient_id")), 1))

# 5. Show results
print("=== Patient Vital Telemetry Analysis ===")
analytical_df.show(truncate=False)

Rendered Output:

=== Patient Vital Telemetry Analysis ===
+----------+------------+----------+---------------+-----------------+--------------+
|patient_id|reading_time|heart_rate|prev_heart_rate|heart_rate_change|patient_avg_hr|
+----------+------------+----------+---------------+-----------------+--------------+
|PAT_01    |10:00       |72        |null           |null             |78.3          |
|PAT_01    |10:05       |78        |72             |6                |78.3          |
|PAT_01    |10:10       |85        |78             |7                |78.3          |
|PAT_02    |10:00       |95        |null           |null             |102.3         |
|PAT_02    |10:05       |110       |95             |15               |102.3         |
|PAT_02    |10:10       |102       |110            |-8               |102.3         |
+----------+------------+----------+---------------+-----------------+--------------+