Window Functions

groupBy turns each group into one row. A window function keeps every row, and adds a value computed over a group of rows. That group is the window.

You know the idea from the SQL chapter Window Functions & Analytics Patterns. This chapter shows the DataFrame form, with Window and .over().

The examples use orders.csv from chapter 3 and the order items.

Window.partitionBy names the column that splits the rows into windows. Here each customer is one window:

from pyspark.sql import functions as F, Window

by_customer = Window.partitionBy("customer_id")

orders.select(
    "order_id", "customer_id", "amount",
    F.sum("amount").over(by_customer).alias("customer_total"),
).orderBy("order_id").show()

orderBy inside the window sorts the rows of each window. Ranking, lag and running totals need it.

Spark must bring the rows of one window together, so a window step shuffles by the partitionBy column (chapter 5).

Three functions number the rows in order. They agree until two rows tie. The order items have a tie: five items have quantity 1.

w = Window.orderBy(F.col("quantity").desc())

order_items.select(
    "order_id", "product", "quantity",
    F.row_number().over(w).alias("row_number"),
    F.rank().over(w).alias("rank"),
    F.dense_rank().over(w).alias("dense_rank"),
).show()

Here rank and dense_rank agree. After a tie, rank skips numbers and dense_rank does not. A sixth item with quantity 0 would get rank 8 but dense_rank 4.

Among the tied rows, row_number picks an order you cannot predict. Add a second sort column, such as order_id, to fix it.

Keep the newest row per key

Chapter 8 said dropDuplicates keeps any one row. To keep the newest order of each customer every time, number the rows newest first and keep number 1:

newest_first = Window.partitionBy("customer_id").orderBy(F.col("order_date").desc())

(orders
 .withColumn("rn", F.row_number().over(newest_first))
 .filter("rn = 1")
 .select("customer_id", "order_id", "order_date", "status")
 .orderBy("customer_id")
 .show())

lag reads a value from the row before, inside the same window. lead reads from the row after:

per_customer = Window.partitionBy("customer_id").orderBy("order_date")

orders.select(
    "customer_id", "order_id", "amount",
    F.lag("amount").over(per_customer).alias("previous"),
).orderBy("customer_id", "order_id").show()

A running total adds up all rows from the start to the current row. The frame says which rows count. rowsBetween(Window.unboundedPreceding, Window.currentRow) means: from the first row up to this one.

running = (Window.orderBy("order_date")
           .rowsBetween(Window.unboundedPreceding, Window.currentRow))

(orders
 .filter(F.col("status") == "paid")
 .select("order_id", "order_date", "amount",
         F.sum("amount").over(running).alias("running_total"))
 .show())

Without rowsBetween, the default frame counts by value, not by row. Two orders on the same day would then share one running total. Give the frame yourself.