Spark SQL window functions | Data Engineering

Q: What are common interview follow-up questions?

TopN per group – use ROW_NUMBER with a filter: df.withColumn("rn", row_number().over(w)).filter(col("rn") <= 3) gives the top 3 transactions per account. Deduplication – use ROW_NUMBER to pick one survivor per business key: partition by the key, order by the tiebreaker, filter where row_number equal

Definition

Window functions compute per-row values over a PARTITION BY and ORDER BY window – ranking, LAG/LEAD, running sums, moving averages – and unlike groupBy they annotate every row instead of collapsing groups into one summary row.

Window function – a function that computes a value for each row based on a window of related rows defined by an OVER clause, without collapsing the rows.
OVER clause – specifies PARTITION BY to divide rows into groups, ORDER BY to order within each group, and optionally a frame bounding which rows contribute.
Ranking functions – ROW_NUMBER (unique sequence), RANK (ties with gaps), DENSE_RANK (ties without gaps).
Offset functions – LAG (read from a preceding row), LEAD (read from a following row).
Aggregate windows – SUM, AVG, COUNT used OVER an ordered window produce running totals and moving averages.
Difference from groupBy – groupBy collapses each group to a single summary row; a window function keeps every original row and annotates it.

At the National Bank of India imagine a statement listing every transaction for an account. A groupBy is like replacing the whole statement with a single line showing only the total – you lose the individual transactions. A window function is like keeping every transaction line but adding extra columns beside each one: its running balance, its rank by amount that month, and the amount of the previous transaction for comparison. The detail stays; you just enrich each line with context computed from its neighbours in the same account, in date order.

Understanding the Concepts

What are window functions in Spark SQL?

Window functions compute a value for each row based on a set of related rows. The key property is that they do not collapse rows. Every original row remains in the output, enriched with the computed value.

PYTHON

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum as spark_sum, row_number, rank, dense_rank
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("WindowDemo").getOrCreate()

# Transaction data at the National Bank of India
txns = spark.createDataFrame([
    ("SAV2001", "2024-01-01", 50000.00),
    ("SAV2001", "2024-01-05", 25000.00),
    ("SAV2001", "2024-01-10", 75000.00),
    ("SAV2001", "2024-01-15", 25000.00),
    ("CUR3001", "2024-01-02", 100000.00),
    ("CUR3001", "2024-01-08", 50000.00),
], ["account_id", "txn_date", "amount"])

How does the OVER clause define the window?

The OVER clause has three parts: PARTITION BY splits data into independent groups, ORDER BY orders rows within each group, and an optional frame clause bounds which rows feed the calculation.

PYTHON

# Define a window partitioned by account, ordered by date
w = Window.partitionBy("account_id").orderBy("txn_date")

# Without a frame, an ordered aggregate defaults to
# ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW (running window)

# Explicit frame for a 3-row moving window
w_moving = Window.partitionBy("account_id").orderBy("txn_date") \
    .rowsBetween(-2, 0)  # current row and 2 preceding

How do ROW_NUMBER, RANK, and DENSE_RANK differ?

All three assign positions within a partition, but they handle ties differently.

| Function | Ties | Gaps After Tie | Use Case |
|---|---|---|---|
| ROW_NUMBER | Different numbers | No ties | Dedup to one row per group |
| RANK | Same rank | Yes (1,1,3) | Ranking with gap awareness |
| DENSE_RANK | Same rank | No (1,1,2) | Consecutive rank numbers |

PYTHON

# Rank transactions by amount within each account
ranked = txns \
    .withColumn("row_num", row_number().over(w)) \
    .withColumn("rnk", rank().over(
        Window.partitionBy("account_id").orderBy(col("amount").desc()))) \
    .withColumn("dense_rnk", dense_rank().over(
        Window.partitionBy("account_id").orderBy(col("amount").desc())))
ranked.show()

TEXT

+----------+----------+--------+-------+---+---------+
|account_id|  txn_date|  amount|row_num|rnk|dense_rnk|
+----------+----------+--------+-------+---+---------+
|   SAV2001|2024-01-10|75000.00|      1|  1|        1|
|   SAV2001|2024-01-01|50000.00|      2|  2|        2|
|   SAV2001|2024-01-05|25000.00|      3|  3|        3|
|   SAV2001|2024-01-15|25000.00|      4|  3|        3|
|   CUR3001|2024-01-02|100000.0|      1|  1|        1|
|   CUR3001|2024-01-08|50000.00|      2|  2|        2|
+----------+----------+--------+-------+---+---------+

How do LAG and LEAD work?

LAG reads a value from a preceding row and LEAD from a following row in the ordered partition. These enable period-over-period comparisons without a self-join.

PYTHON

from pyspark.sql.functions import lag, lead

w_date = Window.partitionBy("account_id").orderBy("txn_date")

comparison = txns \
    .withColumn("prev_amount", lag("amount", 1).over(w_date)) \
    .withColumn("next_amount", lead("amount", 1).over(w_date)) \
    .withColumn("change", col("amount") - col("prev_amount"))
comparison.show()

TEXT

+----------+----------+--------+-----------+-----------+---------+
|account_id|  txn_date|  amount|prev_amount|next_amount|   change|
+----------+----------+--------+-----------+-----------+---------+
|   SAV2001|2024-01-01|50000.00|       null|   25000.00|     null|
|   SAV2001|2024-01-05|25000.00|   50000.00|   75000.00|-25000.00|
|   SAV2001|2024-01-10|75000.00|   25000.00|   25000.00| 50000.00|
|   SAV2001|2024-01-15|25000.00|   75000.00|       null|-50000.00|
+----------+----------+--------+-----------+-----------+---------+

How do running sums and moving averages work?

Aggregate functions used OVER an ordered window produce cumulative and moving results. The same SUM that collapses rows under groupBy instead annotates each row under OVER.

PYTHON

from pyspark.sql.functions import avg as spark_avg

# Running balance per account (cumulative sum)
balance = txns.withColumn("running_balance",
    spark_sum("amount").over(w))
balance.show()

# 3-transaction moving average
moving = txns.withColumn("moving_avg_3",
    spark_avg("amount").over(w_moving))
moving.show()

TEXT

Running Balance:
+----------+----------+--------+---------------+
|account_id|  txn_date|  amount|running_balance|
+----------+----------+--------+---------------+
|   SAV2001|2024-01-01|50000.00|       50000.00|
|   SAV2001|2024-01-05|25000.00|       75000.00|
|   SAV2001|2024-01-10|75000.00|      150000.00|
|   SAV2001|2024-01-15|25000.00|      175000.00|
+----------+----------+--------+---------------+

How do window functions differ from groupBy?

This is the most important distinction. A groupBy collapses each group to a single summary row, losing individual records. A window function keeps every row and adds the computed value alongside.

PYTHON

# groupBy: collapses to one row per account
grouped = txns.groupBy("account_id").agg(spark_sum("amount").alias("total"))
grouped.show()
# Output: 2 rows (one per account)

# Window: keeps all 6 rows, adds running total
windowed = txns.withColumn("running_total", spark_sum("amount").over(w))
windowed.show()
# Output: 6 rows (every transaction preserved)

What about performance and skew with window functions?

Window functions require a shuffle to bring each partition's rows together and a sort to order them within the partition. A window partitioned by a skewed key suffers the same straggler-task problem as a skewed join.

PYTHON

# If one account has 10M transactions, the window on that account
# runs as a single task processing all 10M rows
# Mitigation: break the window into sub-windows where possible,
# or pre-aggregate before windowing

# Also: specify ORDER BY precisely to avoid non-deterministic ROW_NUMBER
# WRONG: ambiguous order makes ROW_NUMBER non-deterministic
w_bad = Window.partitionBy("account_id")  # no ORDER BY!
# ROW_NUMBER without ORDER BY assigns arbitrary numbers

# CORRECT: always specify ORDER BY for ranking functions
w_good = Window.partitionBy("account_id").orderBy("txn_date", "amount")

What are common interview follow-up questions?

DIAGRAM

- **Top-N per group** – use ROW_NUMBER with a filter: `df.withColumn("rn", row_number().over(w)).filter(col("rn") <= 3)` gives the top 3 transactions per account.
- **Deduplication** – use ROW_NUMBER to pick one survivor per business key: partition by the key, order by the tiebreaker, filter where row_number equals 1.
- **Gap detection** – use LAG to compare each transaction date to the previous one and flag gaps exceeding a threshold.

PYTHON

# Top-N per group: top 2 transactions per account by amount
w_top = Window.partitionBy("account_id").orderBy(col("amount").desc())
top2 = txns.withColumn("rn", row_number().over(w_top)) \
    .filter(col("rn") <= 2) \
    .drop("rn")
top2.show()

Visual Overview

DIAGRAM

  groupBy (collapses rows)
  ┌──────────────────────────────┐
  │ SAV2001  50000               │
  │ SAV2001  25000               │  groupBy(account_id)
  │ SAV2001  75000               │  .sum(amount)
  │ SAV2001  25000               │
  └──────────────┬───────────────┘
                 ▼
  ┌──────────────────────────────┐
  │ SAV2001  175000  (1 row)     │  Detail lost!
  └──────────────────────────────┘

  Window Function (annotates rows)
  ┌──────────────────────────────┐
  │ SAV2001  50000               │  SUM(amount)
  │ SAV2001  25000               │  OVER(PARTITION BY account_id
  │ SAV2001  75000               │       ORDER BY txn_date)
  │ SAV2001  25000               │
  └──────────────┬───────────────┘
                 ▼
  ┌──────────────────────────────────────────┐
  │ SAV2001  50000   running_bal=  50000     │
  │ SAV2001  25000   running_bal=  75000     │  All 4 rows kept!
  │ SAV2001  75000   running_bal= 150000     │
  │ SAV2001  25000   running_bal= 175000     │
  └──────────────────────────────────────────┘

Key Takeaway

Window functions in Spark SQL compute a value for each row based on a window of related rows defined by an OVER clause, and the defining property is that they do not collapse rows. The OVER clause has a PARTITION BY that splits rows into independent groups, an ORDER BY that orders rows within each group, and optionally a frame that bounds which rows feed the calculation.

Ranking functions assign positions within a partition: ROW_NUMBER gives a unique increasing number so ties differ, RANK ties at the same number then skips, and DENSE_RANK ties without skipping. Offset functions LAG and LEAD read a value from an earlier or later row, enabling period-over-period comparisons in one pass without self-joins. Aggregate functions like SUM and AVG used OVER an ordered window give running totals and moving averages while preserving every original row.

The contrast with groupBy is the heart of it. A groupBy collapses each group into one summary row, losing individual records. A window function keeps every row and annotates it with the computed value, so a statement of a hundred transactions stays a hundred rows, each carrying its running balance, rank, or previous amount. At the National Bank of India, running balance per account was SUM over a window partitioned by account and ordered by timestamp, and comparing each transaction to the previous one was LAG.

The nuances for a senior engineer are that window functions require a shuffle and sort, so a skewed partition key causes the same straggler problems as a join, and that ORDER BY and frame must be specified precisely because an ambiguous order makes ROW_NUMBER non-deterministic and a wrong frame silently produces incorrect running or moving values.

Remember one sentence: window functions annotate every row with context from its neighbours without collapsing the group, which is the opposite of groupBy.

---