Spark SQL - Advanced Transformations: Windowing Hands-on Workbook

This workbook walks you through tracing transactions using Window functions.

1. The Dataset

Assume we have a customer transactions dataset:

cust_id,tx_date,amount
1,2026-05-01,100
2,2026-05-01,150
1,2026-05-02,50
1,2026-05-03,200
2,2026-05-02,300

2. Tasks

Task 1: Trace Running Total

Write the PySpark DataFrame code to calculate a running total of the transaction amount per cust_id ordered by tx_date. Trace step-by-step partition aggregates.

Task 2: Custom Rank Extraction

Write the PySpark code to rank transactions for each customer by amount in descending order. Extract only the highest transaction for each customer.

3. Step-by-Step Solutions

Solution 1: Tracing Running Totals

• PySpark Code:

from pyspark.sql import Window
from pyspark.sql.functions import col, sum as _sum

windowSpec = Window.partitionBy("cust_id").orderBy("tx_date")
result_df = df.withColumn("running_total", _sum("amount").over(windowSpec))
result_df.show()

• Partition-Level Tracing:
  • Group cust_id = 1:
    • Sorted Records:
      1. 2026-05-01, amount=100 running_total = 100
      2. 2026-05-02, amount=50 running_total = 100 + 50 = 150
      3. 2026-05-03, amount=200 running_total = 150 + 200 = 350
  • Group cust_id = 2:
    • Sorted Records:
      1. 2026-05-01, amount=150 running_total = 150
      2. 2026-05-02, amount=300 running_total = 150 + 300 = 450
• Final Output:

+-------+----------+------+-------------+
|cust_id|   tx_date|amount|running_total|
+-------+----------+------+-------------+
|      1|2026-05-01|   100|          100|
|      1|2026-05-02|    50|          150|
|      1|2026-05-03|   200|          350|
|      2|2026-05-01|   150|          150|
|      2|2026-05-02|   300|          450|
+-------+----------+------+-------------+

Solution 2: Tracing Ranks

• PySpark Code:

from pyspark.sql.functions import rank

rankSpec = Window.partitionBy("cust_id").orderBy(col("amount").desc())
ranked_df = df.withColumn("tx_rank", rank().over(rankSpec))
top_tx = ranked_df.filter(col("tx_rank") == 1)
top_tx.show()

• Output Trace:
  • Group 1 sorted: [(200, rank=1), (100, rank=2), (50, rank=3)] Keeps (200)
  • Group 2 sorted: [(300, rank=1), (150, rank=2)] Keeps (300)
• Final Table:

+-------+----------+------+-------+
|cust_id|   tx_date|amount|tx_rank|
+-------+----------+------+-------+
|      1|2026-05-03|   200|      1|
|      2|2026-05-02|   300|      1|
+-------+----------+------+-------+