When you are looking to aggregate your data, the standard PySpark groupBy() function can do all that for you. It’s what it was built for, but it has a fundamental restriction. It only ever returns one row per collection of data records. You SUM a thousand rows, or a million rows, you get one row back.
Often, that’s exactly what you want, but sometimes it would be handy to also get back some additional data from some or all of the rows that went into the aggregation.
That’s where the PySpark Window functions come into play. They let you calculate values across related records without collapsing those records into a single result. This means you keep the transaction details while gaining aggregation information about the wider group.
They are useful when you need rankings, running totals, comparisons with previous records, or calculations within groups.
In this article, I’ll explain how PySpark window functions work and show how to use them for several common tasks:
Ranking rows within groups
Calculating running totals
Comparing current and previous values
Finding each row’s share of a group total
Selecting the top records from each group
My examples use a small sales dataset, but the same techniques apply equally well to larger data sets like event logs, financial records, customer activity, sensor readings, and many other types of ordered or grouped data.
Setting up PySpark
If PySpark is not already installed on your system, create a project folder and add it with uv:
The local[*] setting runs Spark locally and allows it to use the available processor cores. You don’t need a cluster to follow my examples. Just make sure the code above runs without errors by saving it to a suitable file and running this command.
spark-submit spark_example.py
Creating our example dataset
The dataset contains sales made by three stores over several days:
Each row represents one transaction. We’ll use window functions to analyse the transactions while keeping every row in the result.
What is a window?
A window defines the set of rows that PySpark should consider when calculating a value for the current row.
Most window specifications contain one or more of these parts:
partitionBy() divides the data into groups.
orderBy() defines the order of rows inside each group.
rowsBetween() or rangeBetween() defines the window frame relative to the current row.
Here is a basic window specification:
store_window = Window.partitionBy("store")
This divides the data by store. Every London transaction belongs to one window partition, every Manchester transaction belongs to another, and every Bristol transaction belongs to a third.
You can use that window with an aggregate function:
Use groupBy() when you want one result row per group.
Use a window function when you want group-level calculations alongside the original rows.
Ranking rows within each group
Ranking is one of the most common uses of window functions. For example, you might want to rank transactions from largest to smallest within each store.
First, define a window that partitions by store and orders by transactions by amount:
Use row_number() when you need exactly one first, second, or third row. Use rank() or dense_rank() when ties should receive equal treatment.
Selecting the top records from each group
One of the most useful ranking patterns is to select the top one or more records from each group. For example, the following code returns the two largest transactions from each store:
The second example returns the two largest sales across the entire dataset. The window-based example returns two sales from each store.
This pattern is useful for questions such as:
What are the five highest-value orders for each customer?
Which three products sell best in each region?
What are the latest two events for each device?
Calculating running totals
A running total adds the current row’s value to the values from earlier rows.
To calculate cumulative sales for each store, define a window that:
1/ Partitions transactions by store. 2/ Orders them by date and transaction ID. 3/ Starts at the first row in the partition and ends at the current row.
The first transaction in each store begins the total. Each later transaction adds its amount to the previous total. Because the window is partitioned by store, the calculation starts again when the store changes.
The second ordering column, transaction_id, makes the ordering deterministic when two transactions share the same date. Without a clear tie-breaker, rows with equal ordering values may not always appear in the order you expect.
Comparing a row with the previous row
The lag() function retrieves a value from an earlier row in the same window. It is useful for measuring changes over time.
First, define a window that orders transactions chronologically within each store:
For the first row in each store, the average uses one transaction. For the second row, it uses two. From the third row onward, it uses the current transaction and the previous two.
This is a row-based window, not a time-based window. If one store makes several sales in a day and another makes one sale per week, each calculation still covers three rows.
The rowsBetween and rangeBetween window frames
Window frames control which rows contribute to a calculation. rowsBetween() uses row positions. This frame example includes the current row and the previous three rows:
.rowsBetween(-3, Window.currentRow)
rangeBetween() uses values from the ordering column. Rows with ordering values inside the specified range are included. For example, if the ordering column contains Unix timestamps measured in seconds, this window covers the current timestamp and the previous seven days:
The distinction between these two functions is important:
Use rowsBetween() when you want a fixed number of records.
Use rangeBetween() when you want records within a value or time range.
Time-based range windows require care because the ordering expression must use a suitable numeric representation and consistent units.
Reusing window specifications
Window specifications do not modify a DataFrame by themselves. They describe how a calculation should group, order, and frame rows. Defining windows once and reusing them makes code easier to read:
Clear names such as store_date_window and running_total_window also make it easier to understand why each calculation behaves as it does.
Performance consideration
Window functions are useful, but they are not free. PySpark may need to move and sort data so that rows with the same partition key are processed together and appear in the required order. That data shuffle can come at a performance cost.
To see if that happens, you can inspect the execution plan with:
analysed_sales.explain("formatted")
Look for exchange and sort operations. These are often necessary for window calculations, but they can become expensive on large datasets.
Several habits help keep window queries manageable:
If one partition key contains far more rows than the others, one task may have substantially more work to do. For example, partitioning by country can be uneven if most records belong to one country.
Choose partition keys that match the calculation, but be aware of how the data is distributed.
Reuse calculated results carefully
If the same windowed DataFrame is used by several later actions, caching may avoid recalculating it:
If several rows share the same date, ordering only by date may leave their relative order unclear. Add a suitable tie-breaker, such as a transaction ID, when the sequence matters.
Expecting a window to reduce rows
Window functions add calculations to rows; they do not normally reduce the number of rows. To keep only the highest-ranked records, add a rank column and filter it afterwards.
Treating row-based windows as time-based windows
rowsBetween(-6, 0) includes seven rows, not seven days. Use a range-based window when the calculation must cover a specific period.
Putting the techniques together
The following example adds several useful measures to each transaction:
Each transaction now carries information about its store, its position in time, and its ranking by value. The original transaction-level detail remains intact.
Summary
PySpark window functions calculate values across related rows while preserving the original records. They are particularly useful when a groupBy() would remove detail that you still need.
When using Windowing functions, the main takeaways are:
Use partitionBy() to define independent groups.
Use orderBy() when the calculation depends on sequence.
Use a window frame to control which nearby rows contribute.
Use ranking functions to compare rows within groups.
Use lag() and lead() to compare records across time.
Use aggregate functions over windows for totals, shares, and moving calculations.
Window functions often require Spark to repartition, sort and shuffle data, so inspect execution plans if processing slows down and reduce the inputs by filtering early where possible. Once you understand how the partition, ordering, and window frames work together, window functions become a practical tool for solving many common data-engineering problems.
