Skip to main content
PySpark intermediate Lesson 6 of 10

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.

Frequently Asked Questions

What is the difference between a window function and groupBy?
groupBy collapses each group to one row. A window function keeps every row and adds a column computed over that row's window, so you can put a group total next to each individual row without a join back.
When do I need rowsBetween?
Whenever the default frame is wrong for you. With an orderBy, the default frame is every row from the start of the partition to the current row, which gives a running total. For a whole-partition total you need rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing).
What is the difference between rank, dense_rank, and row_number?
On ties, rank leaves gaps (1, 2, 2, 4), dense_rank does not (1, 2, 2, 3), and row_number breaks ties arbitrarily so every row gets a distinct number (1, 2, 3, 4). Use row_number for deduplication, dense_rank for ranking categories, rank when gaps are meaningful.
Why is my window function slow?
A window with partitionBy shuffles on that key, exactly like groupBy. A window with no partitionBy is far worse — it moves the entire DataFrame into a single partition on one executor, which will not scale and may run out of memory.