A Practical Introduction to PySpark Window Functions

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 groupsCalculating running totalsComparing current and previous valuesFinding each row’s share of a group totalSelecting the top records from each groupMy 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.Table of contentsSetting up PySparkCreating our example datasetWhat is a window?Ranking rows within each grouprow_number, rank, and dense_rankSelecting the top records from each groupCalculating running totalsComparing a row with the previous rowCalculating a row’s share of a group totalCalculating moving averagesThe rowsBetween and rangeBetween window framesReusing window specificationsPerformance considerationFilter earlySelect only the required columnsWatch for skewed partitionsReuse calculated results carefullyCommon mistakesForgetting partitionByUsing an incomplete orderingExpecting a window to reduce rowsTreating row-based windows as time-based windowsPutting the techniques togetherSummarySetting up PySparkIf PySpark is not already installed on your system, create a project folder and add it with uv:Next, we can test whether the installation worked OK by creating a Spark session: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. Creating our example datasetThe 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: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:The result still contains every transaction, but each row now includes the total sales for its store.In contrast, a groupBy() would produce only three rows:The difference is important: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 groupRanking 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:Then apply row_number():The largest sale in each store receives rank 1, the second-largest receives rank 2, and so on.row_number, rank, and dense_rankPySpark provides three closely related ranking functions:They differ when two rows have the same ordering value:row_number() always assigns a unique sequential numberrank() gives tied rows the same rank and leaves gaps afterwards.dense_rank() gives tied rows the same rank without leaving gaps.Imagine sales amounts of 100, 100, and 80. The results for the three rankings would be: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 groupOne 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:This is different from, say,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 totalsA 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.Now we can apply the sum() aggregation to the above window: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 rowThe 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:Now we can use the lag() function on that window to retrieve the previous transaction amount:The first transaction for each store has no previous transaction, so the previous_amount and change_from_previous columns are NULL.You can also compare dates.This can help identify gaps in customer activity, delayed events, or changes in daily measurements.The related lead() function looks forward instead of backwards.Window aggregates make it easy to compare each row with its group.The following code calculates each transaction’s percentage of its store’s total sales:Because the store total appears alongside every transaction, there is no need to aggregate the data and join the totals back to the original rows.This same pattern can calculate:An employee’s salary as a share of departmental payrollA product’s sales as a share of its categoryA transaction’s value compared with a customer’s total spendingCalculating moving averagesA running total includes every earlier row in the partition. A moving calculation uses a limited number of nearby rows.For example, the following window includes the current transaction and the previous two transactions:We can now use this window to calculate a three-transaction moving average: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 framesWindow frames control which rows contribute to a calculation.rowsBetween() uses row positions. This frame example includes the current row and the previous three rows: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 specificationsWindow 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:You can then apply several calculations: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 considerationWindow 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: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:Filter earlyRemove unnecessary rows before applying a window:Filtering first reduces the amount of data that Spark may need to move and sort.Select only the required columnsIf the calculation needs only a few columns, remove the others before the window operation:Watch for skewed partitionsIf 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 carefullyIf the same windowed DataFrame is used by several later actions, caching may avoid recalculating it:Caching only helps when the result is reused. Do not automatically cache every intermediate DataFrame.Common mistakesForgetting partitionByThis window ranks every row across the complete dataset:That may be correct, but it is not the same as ranking sales separately within each store:Using an incomplete orderingIf 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 rowsWindow 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 windowsrowsBetween(-6, 0) includes seven rows, not seven days. Use a range-based window when the calculation must cover a specific period.Putting the techniques togetherThe 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.SummaryPySpark 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.