How to Use Window Functions in PySpark: Step by Step
Window functions in PySpark let you number, order and compare rows within groups without losing the detail of each record. Unlike groupBy, which collapses everything into a single row per group, a window function keeps every row and adds the calculation in a new column. This is how you build rankings, running totals and comparisons with the previous row. Imagine you want the top salespeople in each region, or a running sales total: these are perfect cases for window functions.
Prerequisites
- Have PySpark installed, or an environment such as Databricks or Microsoft Fabric already set up.
- Know how to create a DataFrame and use
withColumn. - Basic Python knowledge.
Step 1: Create sample data
We start by creating a SparkSession and a small DataFrame with sales by person and by region. A small example makes it easy to see the effect of each function. Keep this code, because we will build on it in the next steps.
from pyspark.sql import SparkSession
spark = SparkSession.builder.appName("janelas").getOrCreate()
dados = [
("Ana", "Norte", 100),
("Bruno", "Norte", 150),
("Carla", "Norte", 120),
("Diogo", "Sul", 200),
("Eva", "Sul", 180),
]
df = spark.createDataFrame(dados, ["nome", "regiao", "vendas"])
df.show()
Step 2: Define the window
A window defines how Spark should look at the rows before calculating. partitionBy splits the data into groups — here, one window per region — and orderBy sets the order within each group, in this case from the highest sales to the lowest. Nothing is calculated yet: we are only describing the window we will reuse.
Separating defining the window from applying it keeps the code readable: we can reuse the same janela across several calculations without repeating it.
from pyspark.sql.window import Window
from pyspark.sql import functions as F
janela = Window.partitionBy("regiao").orderBy(F.col("vendas").desc())
Step 3: Number rows with row_number()
The row_number() function assigns a sequential position — 1, 2, 3... — within each region, following the window order. It is the most direct way to create a ranking. We use withColumn to store the result in a new column called posicao.
df = df.withColumn("posicao", F.row_number().over(janela))
df.show()
Step 4: Calculate a running total
When we apply sum() to an ordered window, Spark does not add up the whole group at once: it adds cumulatively, row by row, up to the current row. This is exactly what we need for a running total within each region.
By default, an ordered window uses the frame from the start of the group up to the current row, and that is what produces the cumulative effect we see in the column.
df = df.withColumn("acumulado", F.sum("vendas").over(janela))
df.show()
Step 5: Compare with the previous row using lag()
The lag() function fetches the value from a previous row within the same window — by default, the row immediately above. It is very useful to calculate variations, such as the sales difference between two positions. The first row of each group has no predecessor, so it gets null.
df = df.withColumn("venda_anterior", F.lag("vendas").over(janela))
df.show()
Check the result
Look at the final result. In the Norte region, Bruno comes first (150), followed by Carla (120) and Ana (100), so the posicao column shows 1, 2 and 3. The acumulado column grows through the group — 150, then 270, then 370 — and restarts in the Sul region. The venda_anterior column is null on the first row of each region, because there is no row before it. If you see these values, the window functions are working correctly.
To try it out, paste each code block in order into a notebook and run df.show() after every step, watching the new columns appear one by one.
Conclusion
With just a few lines of code, you created a ranking, a running total and a row-to-row comparison, all without losing the detail of the original data. From here, try rank() and dense_rank() to handle ties, or lead() to look at the next row. For moving averages, adjust the window frame with rowsBetween. Which calculation in your daily work would become simpler with a window function?