DEV Community

Karan Pulyala
Karan Pulyala

Posted on

PySpark Window Functions on Databricks: Latest-Record Dedup, Running Totals, Sessionisation, and the Frame Traps

Window functions show up in almost every PySpark interview, and in almost every real pipeline: deduplicating CDC feeds, picking the latest record per key, running totals, sessionisation. They're also an easy place to write code that's correct on a sample and painfully slow (or subtly wrong) on production data.

This post walks through the patterns I use most, with code you can paste into a Databricks notebook, and the gotchas interviewers like to probe.

The setup

from pyspark.sql import functions as F
from pyspark.sql.window import Window

events = spark.createDataFrame(
    [
        ("c1", "2026-10-01 09:00:00", "view",     0.0),
        ("c1", "2026-10-01 09:05:00", "cart",     0.0),
        ("c1", "2026-10-01 09:07:00", "purchase", 40.0),
        ("c1", "2026-10-01 11:30:00", "view",     0.0),
        ("c2", "2026-10-01 10:00:00", "view",     0.0),
        ("c2", "2026-10-01 10:00:00", "view",     0.0),   # exact duplicate
        ("c2", "2026-10-02 08:15:00", "purchase", 15.0),
    ],
    "customer_id string, event_ts string, event_type string, amount double",
).withColumn("event_ts", F.to_timestamp("event_ts"))
Enter fullscreen mode Exit fullscreen mode

A window spec has three parts: partitionBy (which rows belong together), orderBy (their order inside the group), and an optional frame (rowsBetween / rangeBetween, which rows around the current one the function can see).

Pattern 1: latest record per key (and why row_number beats rank here)

w_latest = Window.partitionBy("customer_id").orderBy(F.col("event_ts").desc())

latest = (
    events
    .withColumn("rn", F.row_number().over(w_latest))
    .filter("rn = 1")
    .drop("rn")
)
Enter fullscreen mode Exit fullscreen mode

The interview question is usually: "What's the difference between row_number, rank, and dense_rank?"

  • row_number always gives 1, 2, 3, and so on, even for ties, so you get exactly one row per key.
  • rank gives tied rows the same number and leaves gaps (1, 1, 3).
  • dense_rank gives ties the same number with no gaps (1, 1, 2).

If two rows tie on event_ts, rank() == 1 returns both and your "latest" table quietly has duplicates. But row_number has its own trap: with ties, which row wins is non-deterministic and can change between runs. Always add a tiebreaker column so the result is reproducible:

w_latest = (
    Window.partitionBy("customer_id")
    .orderBy(F.col("event_ts").desc(), F.col("event_type").desc())
)
Enter fullscreen mode Exit fullscreen mode

In a CDC feed, the tiebreaker is usually a sequence number or commit version, not a business column.

Pattern 2: dedup vs dropDuplicates

For exact duplicates, dropDuplicates(["customer_id", "event_ts", "event_type"]) is simpler and usually cheaper. Reach for row_number when you need to control which duplicate survives (newest ingested_at, highest version). Interviewers like asking which row dropDuplicates keeps. The honest answer is "an arbitrary one", so don't rely on it when the choice matters.

Pattern 3: running totals and the default frame trap

w_running = (
    Window.partitionBy("customer_id")
    .orderBy("event_ts")
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)
)

running = events.withColumn("running_spend", F.sum("amount").over(w_running))
Enter fullscreen mode Exit fullscreen mode

Why spell out rowsBetween? When a window has an orderBy but no explicit frame, Spark defaults to RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW. With RANGE, all rows that share the same event_ts are treated as peers and included together. For customer c2, the two duplicate 10:00 rows would both show the combined total instead of stepping up one at a time. ROWS counts physical rows, which is what most people actually mean by "running total".

The same default bites last():

# Looks right, but returns the current row's value with the default frame
F.last("event_type").over(Window.partitionBy("customer_id").orderBy("event_ts"))

# What you probably wanted: last value in the whole partition
w_all = (
    Window.partitionBy("customer_id")
    .orderBy("event_ts")
    .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
)
F.last("event_type").over(w_all)
Enter fullscreen mode Exit fullscreen mode

Pattern 4: time-based rolling windows with rangeBetween

"Spend in the last 24 hours" is a time range, not a row count. rangeBetween needs a numeric order column, so cast the timestamp to seconds:

w_24h = (
    Window.partitionBy("customer_id")
    .orderBy(F.col("event_ts").cast("long"))
    .rangeBetween(-24 * 3600, 0)
)

rolling = events.withColumn("spend_24h", F.sum("amount").over(w_24h))
Enter fullscreen mode Exit fullscreen mode

rowsBetween(-6, 0) would give you "the last seven events", which is a different question. Mixing those up is a classic interview mistake.

Pattern 5: lag/lead and sessionisation (gaps and islands)

Split each customer's events into sessions whenever there's a gap of more than 30 minutes:

w_order = Window.partitionBy("customer_id").orderBy("event_ts")

sessions = (
    events
    .withColumn("prev_ts", F.lag("event_ts").over(w_order))
    .withColumn(
        "gap_min",
        (F.col("event_ts").cast("long") - F.col("prev_ts").cast("long")) / 60,
    )
    .withColumn(
        "new_session",
        F.when(F.col("prev_ts").isNull() | (F.col("gap_min") > 30), 1).otherwise(0),
    )
    .withColumn(
        "session_no",
        F.sum("new_session").over(
            w_order.rowsBetween(Window.unboundedPreceding, Window.currentRow)
        ),
    )
    .withColumn("session_id", F.concat_ws("-", "customer_id", "session_no"))
)
Enter fullscreen mode Exit fullscreen mode

The trick is the "flag, then cumulative sum" pattern: mark the first row of each island with 1, and a running sum of those flags gives every row its island number. The same idea solves "consecutive login days", "status change periods", and most other gaps-and-islands questions.

For customer c1, the 09:00, 09:05, and 09:07 events land in session 1 and the 11:30 view starts session 2.

Pattern 6: compare a row with its group without a join

w_cust = Window.partitionBy("customer_id")

with_share = (
    events
    .withColumn("total_spend", F.sum("amount").over(w_cust))
    .withColumn(
        "pct_of_customer",
        F.when(F.col("total_spend") > 0, F.col("amount") / F.col("total_spend")),
    )
)
Enter fullscreen mode Exit fullscreen mode

This replaces the "aggregate, then join back" pattern. It's usually cleaner, and it saves a second scan, but it still shuffles by customer_id.

Performance: what windows really cost

Every distinct partitionBy key set needs a shuffle and a sort. A few practical rules:

  1. Reuse the same window spec. Several functions over identical partitionBy/orderBy are evaluated in one Window operator. Five slightly different specs can mean several shuffles. Check with df.explain() and count the Exchange and Window nodes.
  2. Never use a window without partitionBy on big data. Window.orderBy("event_ts") alone moves every row into a single partition, and Spark logs a warning to tell you so. A global row_number over a billion rows is one task doing all the work.
  3. Watch for skew. A single hot key (a bot account, a null customer ID) becomes one giant sorted partition. AQE's skew join handling helps joins, not windows, so you have to fix this yourself, for example by filtering or handling nulls separately before the window.
  4. For top-1 per key, consider an aggregate. max_by (Spark 3.3+) avoids the sort:
latest_type = events.groupBy("customer_id").agg(
    F.max_by("event_type", "event_ts").alias("latest_event_type")
)
Enter fullscreen mode Exit fullscreen mode

Newer Spark versions can also optimise a row_number() ... = 1 filter with a group limit, but max_by is a clean answer when you only need a few columns.

Interview questions worth practising

  1. Your "latest record per key" job returns duplicates on some days. What would you check first?
  2. Why might a running total show the same value on two adjacent rows?
  3. What's the difference between rowsBetween(-6, 0) and rangeBetween(-6 * 86400, 0)?
  4. How would you find customers who logged in on three or more consecutive days?
  5. A window job has one task that runs for 40 minutes while the rest finish in seconds. What's going on and how would you fix it?

If you can answer those with code, window functions won't trip you up.


I've put together a free 10-question sample of Azure Databricks and PySpark interview questions with worked answers, if you'd like more practice like this: free sample. The full set is in the Azure Databricks & PySpark Interview Pack.

Top comments (0)