Master PySpark Window Functions: ROW_NUMBER, RANK, LAG, LEAD, Totals

If you work with Apache Spark, PySpark, Databricks, or SQL, sooner or later you will encounter a problem that looks deceptively simple and uses Window Functions.

“For every customer, find their latest transaction.”

Or:

“Calculate the running total for each account.”

Or:

“Find the previous month’s revenue for every customer.”

Or:

“Get the top 3 products in every category.”

These problems are difficult to solve efficiently if you think only in terms of groupBy() and aggregations.

This is where Window Functions become extremely useful.

Window functions allow you to perform calculations across a group of related rows without collapsing those rows into a single record.

Apache Spark Learning Path

APACHE SPARK LEARNING PATH

Apache Spark Learning Path

APACHE SPARK LEARNING PATH

✓ 1. What is Apache Spark?
✓ 2. Why Spark is Faster than Hadoop
✓ 3. Spark Architecture
✓ 4. Driver vs Executor
✓ 5. Cluster Managers
→ 6. RDD vs DataFrame vs Dataset
→ 7. SparkSession
✓ 8. Reading Data
✓ 9. Lazy Evaluation
✓ 10. DAG
✓ 11. Stages and Tasks
✓ 12. Shuffle
✓ 13. Narrow vs Wide Transformations
→ 14. Writing Data
→ 15. Data Transformations
→ 16. Data Skew

In this tutorial, we’ll understand window functions from the ground up using PySpark, including:

  • Window.partitionBy()
  • Window.orderBy()
  • row_number()
  • rank()
  • dense_rank()
  • lag()
  • lead()
  • Running totals
  • Moving averages
  • Window frames
  • rowsBetween()
  • rangeBetween()
  • Real-world use cases
  • Performance considerations
  • Common mistakes
  • Interview questions

Table of Contents

  1. What Is a Window Function?
  2. Why Do We Need Window Functions?
  3. Window Functions vs GROUP BY
  4. Understanding the Window Specification
  5. Creating a Window in PySpark
  6. ROW_NUMBER
  7. RANK
  8. DENSE_RANK
  9. RANK vs DENSE_RANK vs ROW_NUMBER
  10. LAG
  11. LEAD
  12. Running Total
  13. Moving Average
  14. Window Frames
  15. ROWS BETWEEN
  16. RANGE BETWEEN
  17. Top N Records Per Group
  18. Finding the Latest Record
  19. Detecting Changes Between Rows
  20. Real-World Example
  21. Performance Considerations in Spark
  22. Common Mistakes
  23. Production Considerations
  24. Interview Questions
  25. Key Takeaways
  26. What’s Next?

1. What Is a Window Function?

A window function performs a calculation across a set of rows related to the current row.

The important part is:

Window functions calculate across multiple rows while keeping the original rows in the result.

Consider this data:

customermonthsales
AliceJan100
AliceFeb150
AliceMar200
BobJan300
BobFeb250

Suppose we want the cumulative sales for each customer.

The result should be:

customermonthsalescumulative_sales
AliceJan100100
AliceFeb150250
AliceMar200450
BobJan300300
BobFeb250550

Notice that we did not reduce Alice’s three rows into one row.

Instead, every row remains available.

That is the fundamental idea behind window functions.


2. Why Do We Need Window Functions?

A normal aggregation such as:

df.groupBy("customer").sum("sales")

produces something like:

customertotal_sales
Alice450
Bob550

The individual transaction or monthly rows disappear.

But many analytical problems require both:

  1. The original row
  2. A calculation involving other rows

For example:

  • Previous transaction
  • Next transaction
  • Customer rank
  • Running total
  • Moving average
  • Top N per category
  • First transaction
  • Last transaction
  • Percentage of group total
  • Difference from previous row

These are ideal use cases for window functions.


3. Window Functions vs GROUP BY

This distinction is one of the most important concepts to understand.

GROUP BY

GROUP BY combines rows into groups and returns fewer rows.

SELECT customer, SUM(sales)
FROM sales
GROUP BY customer;

Window Function

A window function calculates over related rows while preserving the rows.

Conceptually:

SUM(sales) OVER (
PARTITION BY customer
)

The result still contains every sales record.

Think of it this way:

GROUP BY
-----------------------
Many rows
One row per group
WINDOW FUNCTION
-----------------------
Many rows
Many rows
+
Additional calculation

This difference is fundamental to understanding analytical SQL and PySpark.


4. Understanding the Window Specification

A window is generally defined using three important concepts:

PARTITION BY
ORDER BY
WINDOW FRAME

In PySpark:

from pyspark.sql.window import Window
window_spec = (
Window
.partitionBy("customer")
.orderBy("month")
)

Let’s break this down.

partitionBy()

Defines the groups of rows that should be considered together.

Window.partitionBy("customer")

means:

Perform the calculation separately for every customer.

orderBy()

Defines the ordering of rows within each partition.

.orderBy("month")

means:

Process each customer’s records in month order.

Window Frame

Defines which rows around the current row should participate in the calculation.

For example:

.rowsBetween(
Window.unboundedPreceding,
Window.currentRow
)

means:

Include everything from the first row of the partition through the current row.


5. Creating a Window in PySpark

Let’s create a sample DataFrame.

from pyspark.sql import SparkSession
spark = SparkSession.builder \
.appName("WindowFunctions") \
.getOrCreate()
data = [
("Alice", "2026-01", 100),
("Alice", "2026-02", 150),
("Alice", "2026-03", 200),
("Bob", "2026-01", 300),
("Bob", "2026-02", 250),
("Bob", "2026-03", 400)
]
columns = ["customer", "month", "sales"]
df = spark.createDataFrame(data, columns)
df.show()

Output:

+--------+-------+-----+
|customer| month|sales|
+--------+-------+-----+
| Alice|2026-01| 100|
| Alice|2026-02| 150|
| Alice|2026-03| 200|
| Bob|2026-01| 300|
| Bob|2026-02| 250|
| Bob|2026-03| 400|
+--------+-------+-----+

Now create the window:

from pyspark.sql.window import Window
window_spec = (
Window
.partitionBy("customer")
.orderBy("month")
)

This tells Spark:

Partition:
Alice
Bob
Within each partition:
Sort by month

6. ROW_NUMBER

row_number() assigns a sequential number to each row within a window.

from pyspark.sql.functions import row_number
result = df.withColumn(
"row_number",
row_number().over(window_spec)
)
result.show()

Output:

+--------+-------+-----+----------+
|customer| month|sales|row_number|
+--------+-------+-----+----------+
| Alice|2026-01| 100| 1|
| Alice|2026-02| 150| 2|
| Alice|2026-03| 200| 3|
| Bob|2026-01| 300| 1|
| Bob|2026-02| 250| 2|
| Bob|2026-03| 400| 3|
+--------+-------+-----+----------+

The numbering starts again for each customer because of:

.partitionBy("customer")

7. RANK

rank() assigns a rank based on the ordering.

Suppose we have:

data = [
("A", 100),
("B", 200),
("C", 200),
("D", 300)
]

Ordering by sales descending:

window_spec = Window.orderBy(df["sales"].desc())

Using:

rank().over(window_spec)

can produce:

sales rank
300 1
200 2
200 2
100 4

Notice the gap.

After two records receive rank 2, the next rank is 4.


8. DENSE_RANK

dense_rank() behaves similarly to rank() but does not leave gaps.

dense_rank().over(window_spec)

Result:

sales dense_rank
300 1
200 2
200 2
100 3

So:

RANK
1
2
2
4
DENSE_RANK
1
2
2
3

This difference is extremely common in SQL and PySpark interviews.


9. RANK vs DENSE_RANK vs ROW_NUMBER

This is one of the most frequently tested window-function concepts.

FunctionDuplicate valuesGaps?Unique number?
ROW_NUMBER()Different numbersNoYes
RANK()Same rankYesNo
DENSE_RANK()Same rankNoNo

Example:

Scores:
100
90
90
80

ROW_NUMBER

1
2
3
4

RANK

1
2
2
4

DENSE_RANK

1
2
2
3

Which one should you use?

Use ROW_NUMBER() when you need exactly one record.

Use RANK() when ranking ties should create gaps.

Use DENSE_RANK() when ties should share a rank without gaps.


10. LAG

lag() allows you to access a previous row.

Suppose we have:

monthsales
Jan100
Feb150
Mar200

We can retrieve the previous month’s sales:

from pyspark.sql.functions import lag
result = df.withColumn(
"previous_sales",
lag("sales", 1).over(window_spec)
)

Result:

monthsalesprevious_sales
Jan100NULL
Feb150100
Mar200150

This is extremely useful for:

  • Month-over-month growth
  • Comparing transactions
  • Detecting changes
  • Time-series analysis
  • Customer behavior analysis

For example:

from pyspark.sql.functions import col
result = result.withColumn(
"sales_change",
col("sales") - col("previous_sales")
)

11. LEAD

lead() works in the opposite direction.

It retrieves a future row.

from pyspark.sql.functions import lead
result = df.withColumn(
"next_sales",
lead("sales", 1).over(window_spec)
)

Result:

monthsalesnext_sales
Jan100150
Feb150200
Mar200NULL

Think:

LAG
← previous row
CURRENT ROW
LEAD
→ next row

12. Running Total

One of the most common applications of window functions is calculating a running total.

Create a window:

running_window = (
Window
.partitionBy("customer")
.orderBy("month")
.rowsBetween(
Window.unboundedPreceding,
Window.currentRow
)
)

Then:

from pyspark.sql.functions import sum
result = df.withColumn(
"running_total",
sum("sales").over(running_window)
)

Result:

customermonthsalesrunning_total
AliceJan100100
AliceFeb150250
AliceMar200450
BobJan300300
BobFeb250550
BobMar400950

The calculation resets for each customer because of:

.partitionBy("customer")

13. Moving Average

Window functions can also calculate moving averages.

For example, suppose we want a three-row moving average.

moving_window = (
Window
.partitionBy("customer")
.orderBy("month")
.rowsBetween(-2, 0)
)

Then:

from pyspark.sql.functions import avg
result = df.withColumn(
"moving_avg",
avg("sales").over(moving_window)
)

Conceptually:

Current row
+
Previous row
+
Two rows before

For:

100
150
200
250

the moving averages become approximately:

100
125
150
200

This technique is useful for:

  • Sales trends
  • Monitoring metrics
  • Time-series analysis
  • Financial data
  • Operational dashboards

14. Window Frames

Window frames are where window functions become significantly more powerful.

A frame defines:

Which rows should be included in the calculation for the current row?

For example:

.rowsBetween(-2, 0)

means:

2 rows before
Current row

While:

.rowsBetween(
Window.unboundedPreceding,
Window.currentRow
)

means:

First row
...
Current row

Spark’s PySpark API exposes both row-based and range-based frame specifications.


15. ROWS BETWEEN

rowsBetween() defines the frame using physical row positions.

Example:

.rowsBetween(-2, 0)

means:

Current row
+ previous 2 rows

Another example:

.rowsBetween(
Window.unboundedPreceding,
Window.currentRow
)

means:

All previous rows
+
Current row

This is commonly used for running totals.


16. RANGE BETWEEN

rangeBetween() defines the frame based on the value of the ordering expression rather than simply the physical row position.

This distinction becomes important when there are duplicate ordering values.

For example, if multiple rows have the same ordering value, a range-based frame can include all rows whose ordering values fall within the specified range.

In production code, choose between:

rowsBetween()

and:

rangeBetween()

based on the business meaning of the window.

Do not treat them as interchangeable.


17. Top N Records Per Group

One of the most useful real-world patterns is:

Find the top 3 products in every category.

Suppose:

categoryproductsales
ElectronicsLaptop900
ElectronicsPhone800
ElectronicsTablet500
ClothingShoes700
ClothingShirt600
ClothingJeans500

Create:

window_spec = (
Window
.partitionBy("category")
.orderBy(col("sales").desc())
)

Then:

result = df.withColumn(
"rank",
row_number().over(window_spec)
)

Filter:

result.filter(col("rank") <= 3)

This gives the top 3 products per category.

This pattern is extremely common in:

  • E-commerce
  • Customer analytics
  • Recommendation systems
  • Financial reporting
  • Data engineering interviews

18. Finding the Latest Record

Another extremely common production requirement:

Find the latest record for every customer.

Suppose a customer can have multiple records.

Create:

window_spec = (
Window
.partitionBy("customer_id")
.orderBy(col("updated_at").desc())
)

Then:

result = (
df
.withColumn(
"rn",
row_number().over(window_spec)
)
.filter(col("rn") == 1)
.drop("rn")
)

This is a very common deduplication pattern.

Instead of using:

groupBy("customer_id").max("updated_at")

we can preserve the entire latest record.

That distinction matters.


19. Detecting Changes Between Rows

Suppose we want to detect when a customer’s status changes.

Example:

customerdatestatus
AJanActive
AFebActive
AMarInactive
AAprInactive
AMayActive

Use lag():

window_spec = (
Window
.partitionBy("customer")
.orderBy("date")
)
result = df.withColumn(
"previous_status",
lag("status").over(window_spec)
)

Then:

result = result.withColumn(
"status_changed",
col("status") != col("previous_status")
)

Now you can identify the rows where the customer’s state changed.

This pattern is useful for:

  • Customer status tracking
  • CDC pipelines
  • Audit history
  • SCD Type 2 processing
  • Monitoring
  • Event streams

20. Real-World Example

Let’s build a more realistic transaction dataset.

data = [
("C001", "2026-01-01", 100),
("C001", "2026-01-05", 250),
("C001", "2026-01-10", 150),
("C002", "2026-01-02", 500),
("C002", "2026-01-07", 200),
("C002", "2026-01-12", 300)
]
df = spark.createDataFrame(
data,
["customer_id", "transaction_date", "amount"]
)

Create the window:

window_spec = (
Window
.partitionBy("customer_id")
.orderBy("transaction_date")
)

Now calculate the previous transaction:

df = df.withColumn(
"previous_amount",
lag("amount").over(window_spec)
)

Calculate the transaction difference:

df = df.withColumn(
"amount_difference",
col("amount") - col("previous_amount")
)

Calculate the running total:

running_window = (
Window
.partitionBy("customer_id")
.orderBy("transaction_date")
.rowsBetween(
Window.unboundedPreceding,
Window.currentRow
)
)
df = df.withColumn(
"running_total",
sum("amount").over(running_window)
)

The resulting dataset contains:

customer_id
transaction_date
amount
previous_amount
amount_difference
running_total

This is a good example of why window functions are powerful.

One DataFrame can contain:

  • Current transaction
  • Previous transaction
  • Difference
  • Cumulative amount

without collapsing the underlying transaction records.


21. Performance Considerations in Spark

Window functions are powerful, but they are not free.

A window often requires Spark to partition and order data according to the window specification.

For example:

Window
.partitionBy("customer_id")
.orderBy("transaction_date")

means Spark needs data organized appropriately for the window operation.

This can involve significant shuffle and sorting work on large datasets.

Therefore, be careful with:

Window.orderBy(...)

without a meaningful partition.

For a very large dataset, a global window can become expensive because the operation may require extensive data movement and sorting.

Practical considerations

1. Choose partition columns carefully

.partitionBy("customer_id")

should represent a meaningful grouping.

2. Avoid unnecessary windows

Don’t calculate five different windows if the same result can be obtained from one window specification.

3. Inspect the execution plan

Use:

df.explain("formatted")

to understand how Spark plans the operation.

4. Watch for data skew

If one customer has millions of records while most customers have only a few hundred, partitioning by customer can produce an imbalanced workload.

5. Be careful with global ordering

This:

Window.orderBy("timestamp")

can be significantly more expensive than a properly partitioned window.


22. Common Mistakes

Mistake 1: Forgetting partitionBy

This:

Window.orderBy("date")

may create a global window when the business problem actually requires calculations per customer.

Always ask:

Should this calculation restart for each entity?

If yes, use:

.partitionBy(...)

Mistake 2: Using GROUP BY When Rows Must Be Preserved

If you need:

original row
+
calculated value

a window function is often more appropriate than groupBy().


Mistake 3: Using RANK When You Need Exactly N Rows

Suppose you need exactly three records per category.

Using:

rank()

can return more than three rows if ties occur.

If you require exactly three rows:

row_number()

is usually the better choice.


Mistake 4: Ignoring Ordering

Functions such as:

lag()
lead()
row_number()

depend heavily on ordering.

This:

Window.partitionBy("customer")

does not tell Spark which record is first or previous.

You generally need:

Window
.partitionBy("customer")
.orderBy("date")

Mistake 5: Not Thinking About Ties

Suppose two transactions have exactly the same timestamp.

Then:

row_number()

may assign different row numbers based on the resulting ordering.

For deterministic results, add an appropriate tie-breaker:

Window
.partitionBy("customer_id")
.orderBy(
col("transaction_date").desc(),
col("transaction_id").desc()
)

This is particularly important in production pipelines.


23. Production Considerations

When using window functions in production Spark pipelines, consider:

Deterministic ordering

Always make the ordering sufficiently specific.

Data skew

Check whether some partition keys contain dramatically more records than others.

Shuffle cost

Window operations involving partitioning and sorting can generate substantial shuffle.

Memory pressure

Large windows can increase resource requirements.

Explain plans

Use:

df.explain("formatted")

to understand the physical plan.

Avoid unnecessary columns

Project only the columns required before expensive operations when possible.

Validate edge cases

Test:

  • Duplicate timestamps
  • NULL values
  • Empty partitions
  • Single-row partitions
  • Large partitions
  • Ties in ranking

Window functions often look correct on a small dataset but expose subtle issues at production scale.


24. Window Functions Cheat Sheet

RequirementFunction
Sequential numberrow_number()
Ranking with gapsrank()
Ranking without gapsdense_rank()
Previous rowlag()
Next rowlead()
Running sumsum().over()
Running averageavg().over()
First valuefirst_value()
Last valuelast_value()
Group-level calculation while preserving rowsWindow aggregation

The PySpark API includes the Window specification and related frame methods such as rowsBetween() and rangeBetween().


25. Interview Questions

If you’re preparing for a PySpark, Databricks, or Data Engineering interview, expect questions such as:

Beginner

1. What is a window function?

A function that performs calculations across related rows while preserving individual rows in the result.

2. What is the difference between GROUP BY and a window function?

GROUP BY reduces rows into groups, while window functions preserve individual rows.

3. What does partitionBy() do?

It defines the groups within which the window calculation is performed.

4. Why do we use orderBy() in a window?

It defines the ordering of rows within each partition.


Intermediate

5. What is the difference between ROW_NUMBER, RANK, and DENSE_RANK?

The key difference is how they handle ties and ranking gaps.

6. What is the difference between LAG and LEAD?

LAG accesses a previous row, while LEAD accesses a subsequent row.

7. How would you find the top 3 employees in every department?

Use:

Window.partitionBy("department").orderBy(
col("salary").desc()
)

and then apply row_number() or an appropriate ranking function.

8. How would you find the latest record for every customer?

Use row_number() with:

.partitionBy("customer_id")
.orderBy(col("timestamp").desc())

and filter for:

row_number == 1

Advanced

9. What is the difference between rowsBetween() and rangeBetween()?

rowsBetween() defines a frame using row positions, while rangeBetween() defines it based on the values of the ordering expression.

10. Why can window functions be expensive in Spark?

Because partitioning and ordering can require data movement and sorting, potentially causing significant shuffle.

11. How would you troubleshoot a slow window operation?

I would inspect the physical execution plan using:

df.explain("formatted")

and investigate partitioning, sorting, shuffle volume, data skew, and unnecessarily large windows.

12. How would you make a latest-record query deterministic when timestamps can be duplicated?

Add a deterministic tie-breaker:

.orderBy(
col("updated_at").desc(),
col("transaction_id").desc()
)

26. Key Takeaways

Window functions are one of the most important tools in PySpark and analytical SQL.

The core mental model is:

PARTITION
Divide related rows
ORDER
Define their sequence
FRAME
Define which rows participate
FUNCTION
Calculate the result

The most important functions to remember are:

ROW_NUMBER
RANK
DENSE_RANK
LAG
LEAD
SUM
AVG
FIRST_VALUE
LAST_VALUE

If you understand:

Window.partitionBy()
Window.orderBy()
Window.rowsBetween()
Window.rangeBetween()

you can solve a large class of real-world data engineering problems.

And more importantly, you move beyond simply knowing PySpark syntax—you start thinking in terms of data relationships, ordering, partitions, frames, and execution cost.


Continue Learning Apache Spark

If you’re following a structured Spark learning path, continue with:

Previous: [RDD vs DataFrame vs Dataset]

You are here: Window Functions in PySpark

Next: Repartition vs Coalesce in Apache Spark

You may also want to read:

  • [What is Apache Spark?]
  • [Spark Architecture]
  • [DAG in Apache Spark]
  • [Spark Stages and Tasks]
  • [Driver vs Executor]
  • [Shuffle in Apache Spark]
  • [Narrow vs Wide Transformations]
  • [Lazy Evaluation in Apache Spark]

Official Documentation

For the complete and current PySpark Window API, see the official Apache Spark documentation.


Final Thought

Window functions are often introduced as a collection of SQL functions.

That’s the wrong way to learn them.

Instead, think about the question:

“For this row, which other rows should I look at?”

Once you can answer that, the rest becomes much easier:

Which group?
partitionBy()
Which order?
orderBy()
Which rows?
rowsBetween() / rangeBetween()
What calculation?
row_number / rank / lag / lead / sum / avg

That mental model is what makes window functions useful—not just memorizing function names.

Leave a Reply

Discover more from Geeky Codes

Subscribe now to keep reading and get access to the full archive.

Continue reading