Spark Window Functions
A window function calculates a value for each row by looking at a group of related rows. A normal groupBy collapses each group into a single row. A window function keeps every row and adds the extra result beside it. Rankings, running totals, and comparisons with the previous row all depend on windows.
The Layman View
Picture a race with several age groups. A judge ranks runners inside each group and keeps the full results list. Each runner still appears on the sheet and gains a new column named rank. A window function creates that column.
Anatomy of a Window
from pyspark.sql.window import Window
from pyspark.sql import functions as F
w = Window.partitionBy("dept").orderBy(F.col("salary").desc())
- partitionBy splits rows into independent groups.
- orderBy sorts the rows inside each group.
- rowsBetween limits how many neighbouring rows each calculation sees.
Sample Data
data = [("Asha","Sales",52000), ("Ben","Sales",48000), ("Cara","Sales",48000),
("Dev","IT",68000), ("Chitra","IT",75000), ("Elena","HR",45000)]
df = spark.createDataFrame(data, ["name","dept","salary"])
Ranking Functions
df.withColumn("row_num", F.row_number().over(w)) \
.withColumn("rank", F.rank().over(w)) \
.withColumn("dense_rank",F.dense_rank().over(w)).show()
dept name salary | row_number rank dense_rank ----------------------------------------------------- Sales Asha 52000 | 1 1 1 Sales Ben 48000 | 2 2 2 Sales Cara 48000 | 3 2 2 <-- tie IT Chitra 75000 | 1 1 1 IT Dev 68000 | 2 2 2 HR Elena 45000 | 1 1 1
- row_number gives a unique number to every row.
- rank gives tied rows the same number and skips the next numbers.
- dense_rank gives tied rows the same number and skips nothing.
Top N Per Group
Combine row_number with a filter to keep the best rows of every group.
top1 = (df.withColumn("rn", F.row_number().over(w))
.filter("rn = 1")
.drop("rn"))
top1.show()
Looking at Neighbouring Rows
The lag function reads the previous row. The lead function reads the next row.
sales = spark.createDataFrame(
[("Jan",100),("Feb",130),("Mar",120),("Apr",160)], ["month","amount"])
wm = Window.orderBy("month_no")
sales = sales.withColumn("month_no", F.expr("array_position(array('Jan','Feb','Mar','Apr'), month)"))
sales.withColumn("previous", F.lag("amount", 1).over(wm)) \
.withColumn("change", F.col("amount") - F.col("previous")).show()
month amount previous change Jan 100 null null Feb 130 100 30 Mar 120 130 -10 Apr 160 120 40
Running Totals and Moving Averages
run = Window.orderBy("month_no").rowsBetween(Window.unboundedPreceding, Window.currentRow)
mov = Window.orderBy("month_no").rowsBetween(-1, 0)
sales.withColumn("running_total", F.sum("amount").over(run)) \
.withColumn("moving_avg", F.avg("amount").over(mov)).show()
Row Mar with frame rowsBetween(-1, 0): Jan 100 Feb 130 <-- included Mar 120 <-- current row, included Apr 160 moving_avg = (130 + 120) / 2 = 125
Performance Notes
- A window with partitionBy shuffles data so each group sits on one machine. Pick a partition column with many distinct values.
- A window without partitionBy pulls all rows into one partition and slows large jobs.
- Reuse one window definition for several columns.
Choosing the Right Function
| Goal | Function |
|---|---|
| Number every row in a group | row_number |
| Rank with gaps after ties | rank |
| Rank without gaps | dense_rank |
| Compare with the previous row | lag |
| Compare with the next row | lead |
| Running total | sum with rowsBetween |
| Share of the group total | sum over a window without orderBy |
Share of Total Example
A window with only partitionBy sums every row of the group. Divide each salary by that sum to find its share.
wd = Window.partitionBy("dept")
df.withColumn("dept_total", F.sum("salary").over(wd)) \
.withColumn("share_pct", F.round(F.col("salary") / F.col("dept_total") * 100, 1)).show()
Asha earns 52000 inside a Sales total of 148000, which gives a share of 35.1 percent. Window functions deliver this result in a single pass and avoid a separate join.