If you're using PySpark and you're still leaning heavily on groupBy, you're probably doing more work than you need to. Window functions offer a different way to look at your data — not just by grouping it into buckets but by giving each row a sense of its surroundings. This lets you calculate things like running totals, row numbers, or lag values without collapsing your data into fewer rows. So, if you're trying to compare each row to the one before it or get the average of the past few entries in a column, you're in the right place.

Window functions don’t just slice the data; they remember the row’s original context and add more meaning to it. This is useful when you're working with logs, time series data, or anything where row order matters.

What Are Window Functions Really?

At a glance, a window function acts across a window — a defined subset of rows related to the current row. You can think of it like a moving frame. The size and shape of the frame are controlled using the Window spec.

Unlike groupBy, which collapses many rows into one, a window function keeps the row but adds context to it. That means you can run calculations across related rows without changing the structure of your DataFrame.

The three most common types of window operations are:

Ranking functions: Add row numbers, rank, or dense rank within partitions.

Analytic functions: Compare values across rows, such as using lead, lag, or nth_value.

Aggregate functions: Compute sums, averages, counts, etc., but do it across a window instead of grouping everything.

Now, let's see how to use them in PySpark.

Step-by-Step Guide to Using Window Functions in PySpark

Before anything else, bring in the necessary PySpark functions. If you’re familiar with the functions module, most of what you need is already there:

python

from pyspark.sql import SparkSession

from pyspark.sql.window import Window

import pyspark.sql.functions as F

Start by creating a SparkSession if you don’t already have one:

python

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

Now let’s create a sample DataFrame to work with. This gives us something practical to apply the window functions on:

python

data = [

("Alice", "2023-01-01", 100),

("Alice", "2023-01-02", 200),

("Alice", "2023-01-03", 300),

("Bob", "2023-01-01", 400),

("Bob", "2023-01-02", 500),

("Bob", "2023-01-03", 600)

]

columns = ["name", "date", "sales"]

df = spark.createDataFrame(data, columns)

df = df.withColumn("date", F.to_date("date"))

This gives a small DataFrame of daily sales per person.

Next, you’ll need to define a window specification. This decides how the window “moves” and what rows are visible to each calculation. The spec can define both partitioning and ordering:

python

windowSpec = Window.partitionBy("name").orderBy("date")

In this example, each person gets their own set of rows, and within those rows, data is ordered by date.

With the window spec defined, you're ready to apply the actual window functions. Let’s start with a running total:

python

df = df.withColumn("running_total", F.sum("sales").over(windowSpec))

This adds up sales so far for each row, ordered by date and scoped to the individual’s data.

You can calculate plenty of other things with the same window. For example, you might want to compare each row to the one before or after:

python

df = df.withColumn("previous_day_sales", F.lag("sales", 1).over(windowSpec))

df = df.withColumn("next_day_sales", F.lead("sales", 1).over(windowSpec))

Or, if you're ranking entries by date within each person’s partition, try:

python

df = df.withColumn("row_number", F.row_number().over(windowSpec))

df = df.withColumn("rank", F.rank().over(windowSpec))

df = df.withColumn("dense_rank", F.dense_rank().over(windowSpec))

Each of these functions creates a new column that keeps the original row intact while adding more detail — whether it’s the order, a comparison, or a position within the group. None of these functions summarize or reduce rows, which makes them ideal for analysis where the row-level context is important.

How to Use These?

Time-Based Calculations

If you’re working with time series, tracking things like rolling averages or time gaps between events is easier with window functions.

python

df = df.withColumn("rolling_avg", F.avg("sales").over(windowSpec.rowsBetween(-2, 0)))

This takes the average of the current row and two previous ones based on the date ordering within the person's data.

Comparing Current vs Previous Entries

This pattern shows up often when you’re trying to detect spikes, drops, or changes over time.

python

df = df.withColumn("sales_diff", F.col("sales") - F.lag("sales", 1).over(windowSpec))

If sales_diff is negative, that means the person sold less than the day before.

Notes on Performance and Behavior

Window functions are powerful, but they’re not free. If your data is large and your partitions are small, you’ll likely see memory overhead and slower execution.

Here are a few tips to keep in mind:

Partition Size: Too small or too large partitions can hurt performance. Aim for balance to reduce shuffling and memory overhead.

Caching: If you're using multiple window functions with the same spec, cache the DataFrame after the first one to avoid recomputation.

Rows vs. Range: Use rowsBetween for row-specific windows (e.g., two previous rows) and rangeBetween when windowing by value ranges.

Null Handling: Nulls affect lead, lag, and rank. PySpark sorts nulls first by default unless specified otherwise using nullsLast.

Partition Cardinality: Avoid using high-cardinality columns (like IDs or timestamps) for partitioning, as this leads to tiny partitions and excessive shuffling.

Execution Plan: Use df.explain() to inspect physical and logical plans. It's helpful to verify window specs are applied efficiently.

Also, remember that null values can affect functions like lead, lag, and rank. If the sort column has nulls, PySpark puts them first by default — unless you tell it otherwise.

Wrapping Up

Window functions in PySpark let you look across rows without collapsing them. They add context whether you're tracking changes over time, computing rankings, or calculating running totals. Once you get used to setting up the window specs and plugging them into the right function, these become everyday tools — especially for time-based or grouped comparisons. The best way to understand them? Try a few out on real data. Start with a small data frame, define a clear partition and order, and then build up from there. Once you get the hang of it, you'll find yourself relying on window functions more than any other feature in PySpark. Stay tuned for more informative guides.