Window Functions in PySpark
Rank rows within groups, compute running totals, and compare each row to its neighbours — without collapsing the DataFrame the way groupBy does.
groupBy answers “what is the total per team”. Window functions answer “how does this row
compare to its team” — without losing the row. That distinction is most of their value.
The data
from pyspark.sql import SparkSession, functions as F
from pyspark.sql.window import Window
spark = SparkSession.builder.appName("windows").master("local[*]").getOrCreate()
df = spark.createDataFrame(
[
("alice", "engineering", 95000, "2019-03-15"),
("bob", "engineering", 87000, "2021-07-01"),
("frank", "engineering", 87000, "2020-02-11"),
("carol", "design", 78000, "2020-01-20"),
("dave", "design", 82000, "2018-11-05"),
("erin", "sales", 67000, "2022-05-30"),
("gina", "sales", 71000, "2021-09-14"),
],
["name", "team", "salary", "hired"],
)
df.show()
+-----+-----------+------+----------+
| name| team|salary| hired|
+-----+-----------+------+----------+
|alice|engineering| 95000|2019-03-15|
| bob|engineering| 87000|2021-07-01|
|frank|engineering| 87000|2020-02-11|
|carol| design| 78000|2020-01-20|
| dave| design| 82000|2018-11-05|
| erin| sales| 67000|2022-05-30|
| gina| sales| 71000|2021-09-14|
+-----+-----------+------+----------+
Note bob and frank earn the same — that tie is what separates the three ranking
functions.
Keeping every row
With groupBy you lose the individual rows:
df.groupBy("team").agg(F.avg("salary").alias("team_avg")).show()
+-----------+-----------------+
| team| team_avg|
+-----------+-----------------+
|engineering|89666.66666666667|
| design| 80000.0|
| sales| 69000.0|
+-----------+-----------------+
To see each person against their team average you would have to join that back. A window does it in one step:
team = Window.partitionBy("team")
df.withColumn("team_avg", F.round(F.avg("salary").over(team), 0)) \
.withColumn("vs_avg", F.col("salary") - F.round(F.avg("salary").over(team), 0)) \
.show()
+-----+-----------+------+----------+--------+-------+
| name| team|salary| hired|team_avg| vs_avg|
+-----+-----------+------+----------+--------+-------+
|carol| design| 78000|2020-01-20| 80000.0|-2000.0|
| dave| design| 82000|2018-11-05| 80000.0| 2000.0|
|alice|engineering| 95000|2019-03-15| 89667.0| 5333.0|
| bob|engineering| 87000|2021-07-01| 89667.0|-2667.0|
|frank|engineering| 87000|2020-02-11| 89667.0|-2667.0|
| erin| sales| 67000|2022-05-30| 69000.0|-2000.0|
| gina| sales| 71000|2021-09-14| 69000.0| 2000.0|
+-----+-----------+------+----------+--------+-------+
Seven rows in, seven rows out, each carrying its team’s average. This is the whole idea.
Ranking, and the tie problem
ranked = Window.partitionBy("team").orderBy(F.desc("salary"))
df.select(
"name", "team", "salary",
F.rank().over(ranked).alias("rank"),
F.dense_rank().over(ranked).alias("dense"),
F.row_number().over(ranked).alias("row_num"),
).show()
+-----+-----------+------+----+-----+-------+
| name| team|salary|rank|dense|row_num|
+-----+-----------+------+----+-----+-------+
| dave| design| 82000| 1| 1| 1|
|carol| design| 78000| 2| 2| 2|
|alice|engineering| 95000| 1| 1| 1|
| bob|engineering| 87000| 2| 2| 2|
|frank|engineering| 87000| 2| 2| 3|
| erin| sales| 71000| 1| 1| 1|
| gina| sales| 67000| 2| 2| 2|
Look at the engineering tie. rank gives both 2 — and would give the next person 4, skipping
3. dense_rank gives both 2 and the next person 3. row_number gives 2 and 3 arbitrarily,
which is non-deterministic across runs unless you add a tiebreaker to the orderBy.
That non-determinism matters:
stable = Window.partitionBy("team").orderBy(F.desc("salary"), F.asc("name"))
df.select("name", "team", F.row_number().over(stable).alias("n")).show(3)
+-----+-----------+---+
| name| team| n|
+-----+-----------+---+
| dave| design| 1|
|carol| design| 2|
|alice|engineering| 1|
+-----+-----------+---+
Always add a deterministic tiebreaker when using row_number for anything that must be
reproducible.
Top N per group
The canonical use of row_number:
top2 = (
df.withColumn("n", F.row_number().over(Window.partitionBy("team").orderBy(F.desc("salary"))))
.filter(F.col("n") <= 2)
.drop("n")
)
top2.show()
+-----+-----------+------+----------+
| name| team|salary| hired|
+-----+-----------+------+----------+
| dave| design| 82000|2018-11-05|
|carol| design| 78000|2020-01-20|
|alice|engineering| 95000|2019-03-15|
| bob|engineering| 87000|2021-07-01|
| erin| sales| 71000|2021-09-14|
| gina| sales| 67000|2022-05-30|
+-----+-----------+------+----------+
You cannot filter on a window function directly — WHERE runs before the window is
computed. Compute it into a column, then filter. That two-step shape is required, not
stylistic.
Deduplication
The same pattern removes duplicates, keeping the newest record per key:
events = spark.createDataFrame(
[("u1", "2026-01-01", "signup"), ("u1", "2026-01-05", "upgrade"),
("u2", "2026-01-02", "signup"), ("u1", "2026-01-03", "login"),
("u2", "2026-01-06", "cancel")],
["user", "ts", "action"],
)
latest = Window.partitionBy("user").orderBy(F.desc("ts"))
events.withColumn("n", F.row_number().over(latest)).filter("n = 1").drop("n").show()
+----+----------+-------+
|user| ts| action|
+----+----------+-------+
| u1|2026-01-05|upgrade|
| u2|2026-01-06| cancel|
+----+----------+-------+
dropDuplicates cannot do this — it keeps an arbitrary row per key, not the newest.
Running totals and frames
Add an orderBy and aggregates become cumulative:
running = Window.partitionBy("team").orderBy("hired")
df.select(
"name", "team", "hired", "salary",
F.sum("salary").over(running).alias("running_total"),
).show()
+-----+-----------+----------+------+-------------+
| name| team| hired|salary|running_total|
+-----+-----------+----------+------+-------------+
| dave| design|2018-11-05| 82000| 82000|
|carol| design|2020-01-20| 78000| 160000|
|alice|engineering|2019-03-15| 95000| 95000|
|frank|engineering|2020-02-11| 87000| 182000|
| bob|engineering|2021-07-01| 87000| 269000|
| gina| sales|2021-09-14| 71000| 71000|
| erin| sales|2022-05-30| 67000| 138000|
+-----+-----------+----------+------+-------------+
Adding orderBy silently changed the frame from “the whole partition” to “start of partition
through current row”. That is the single most surprising behaviour in window functions —
avg over the same window is a moving average, not the team average, the moment you add
an order.
Make the frame explicit when it matters:
whole = Window.partitionBy("team").orderBy("hired") \
.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
moving = Window.partitionBy("team").orderBy("hired").rowsBetween(-1, 1)
df.select(
"name", "team", "salary",
F.sum("salary").over(whole).alias("team_total"),
F.round(F.avg("salary").over(moving), 0).alias("moving_avg_3"),
).show()
+-----+-----------+------+----------+------------+
| name| team|salary|team_total|moving_avg_3|
+-----+-----------+------+----------+------------+
| dave| design| 82000| 160000| 80000.0|
|carol| design| 78000| 160000| 80000.0|
|alice|engineering| 95000| 269000| 91000.0|
|frank|engineering| 87000| 269000| 89667.0|
| bob|engineering| 87000| 269000| 87000.0|
| gina| sales| 71000| 138000| 69000.0|
| erin| sales| 67000| 138000| 69000.0|
+-----+-----------+------+----------+------------+
Comparing to neighbours
lag and lead reach backwards and forwards:
prices = spark.createDataFrame(
[("2026-01-01", 100.0), ("2026-01-02", 104.0), ("2026-01-03", 99.0),
("2026-01-04", 108.0), ("2026-01-05", 112.0)],
["day", "price"],
)
w = Window.orderBy("day")
prices.select(
"day", "price",
F.lag("price", 1).over(w).alias("prev"),
F.round(((F.col("price") - F.lag("price", 1).over(w)) / F.lag("price", 1).over(w)) * 100, 2).alias("pct_change"),
F.lead("price", 1).over(w).alias("next"),
).show()
+----------+-----+-----+----------+-----+
| day|price| prev|pct_change| next|
+----------+-----+-----+----------+-----+
|2026-01-01|100.0| NULL| NULL|104.0|
|2026-01-02|104.0|100.0| 4.0| 99.0|
|2026-01-03| 99.0|104.0| -4.81|108.0|
|2026-01-04|108.0| 99.0| 9.09|112.0|
|2026-01-05|112.0|108.0| 3.70| NULL|
+----------+-----+-----+----------+-----+
The nulls at the edges are correct — the first row has no previous, the last no next. Supply
a default if you need one: F.lag("price", 1, 0.0).
The no-partition warning
This window has no partitionBy:
prices.select("day", F.sum("price").over(Window.orderBy("day"))).explain()
== Physical Plan ==
Window [sum(price#1) windowspecdefinition(day#0 ASC NULLS FIRST, ...)]
+- *(1) Sort [day#0 ASC NULLS FIRST], false, 0
+- Exchange SinglePartition, ENSURE_REQUIREMENTS
Exchange SinglePartition — every row moved to one executor. On five rows that is fine; on a
billion it will not finish. Spark warns at runtime too:
WARN WindowExec: No Partition Defined for Window operation!
Moving all data to a single partition, this can cause serious performance degradation.
Take that warning seriously. If a global ordering is genuinely required, aggregate first to shrink the data, then window over the smaller result.
Practice
1. Find the highest-paid person per team using a window, then using groupBy. Which is clearer?
df.withColumn("n", F.row_number().over(Window.partitionBy("team").orderBy(F.desc("salary")))) \
.filter("n = 1").select("team", "name", "salary").show()
+-----------+-----+------+
| team| name|salary|
+-----------+-----+------+
| design| dave| 82000|
|engineering|alice| 95000|
| sales| gina| 71000|
+-----------+-----+------+
The groupBy version needs max(struct(...)) or a join back. The window version reads as
what you actually meant, and extends to top-3 by changing one number.
2. Add orderBy to a window and compare avg before and after.
without orderBy: 89667 89667 89667 (team average, all rows equal)
with orderBy: 95000 91000 89667 (running average)
Adding an order changes the default frame to unbounded-preceding-through-current-row. If you
wanted the team average, either drop the orderBy or set the frame explicitly.
3. Compute a 7-day moving average on daily data with fewer than 7 days at the start.
w = Window.orderBy("day").rowsBetween(-6, 0)
prices.select("day", "price", F.round(F.avg("price").over(w), 2).alias("ma7")).show()
+----------+-----+------+
| day|price| ma7|
+----------+-----+------+
|2026-01-01|100.0| 100.0|
|2026-01-02|104.0| 102.0|
|2026-01-03| 99.0| 101.0|
|2026-01-04|108.0|102.75|
|2026-01-05|112.0| 104.6|
+----------+-----+------+
Spark averages whatever rows exist rather than returning null, so early values are averages of fewer points. If a partial window should be null, count the rows in the frame and blank the result where the count is below 7.
4. Use rangeBetween instead of rowsBetween on the salary column. What changes?
rowsBetween counts rows; rangeBetween counts values. With rangeBetween(-1000, 1000)
ordered by salary, bob and frank — both on 87000 — share an identical frame including each
other, whereas rowsBetween(-1, 1) would give them different frames based on position. Use
rangeBetween when the offset means a quantity (within £1000, within 7 days), rowsBetween
when it means a count of records.
Next: running your own Python inside Spark, and what it costs.