Window functions in Spark: practical tips and trade-offs

Window functions in Spark: practical tips and trade-offs

If you've ever needed to rank transactions within a customer, compute a running balance, or find the most recent record per user, you've encountered a problem that window functions solve elegantly. Window functions in Spark let us perform calculations across ordered subsets of data without collapsing rows, which is powerful but comes with performance considerations we need to understand.

What is a window?

A window is a frame of rows that Spark evaluates together. It's defined by two components: partitioning (which rows belong in the same group) and ordering (the sequence within each group). Think of it like sorting and grouping data, then applying a function to each group without losing the original rows.

In a real e-commerce scenario, imagine we have order events and we want to flag the latest order per customer. We partition by user_id so each customer's data stays separate, then order by updated_at descending to get the most recent first.

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# Define a window: partition by user, order by date descending
w = Window.partitionBy("user_id").orderBy(F.col("updated_at").desc())

# Add row number within each partition
df_with_rn = df.withColumn("rn", F.row_number().over(w))

# Keep only the latest record per user
latest_orders = df_with_rn.filter("rn = 1").drop("rn")

latest_orders.show()

This pattern is common in production systems: deduplicate records, fetch the current state, or extract snapshots. The key insight is that the original row count doesn't change until we filter; the window function adds metadata (like rn) alongside existing columns.

Ranking: row_number, rank, and dense_rank

Ranking functions assign a position to rows within a window. The choice between them matters for correctness and business logic.

row_number() assigns a unique integer to every row, even if values are tied. rank() assigns the same rank to ties but skips numbers after (1, 1, 3). dense_rank() assigns the same rank to ties and does not skip numbers (1, 1, 2).

Consider a trading system where we rank orders by execution price within each symbol:

trades_data = [
    ("AAPL", "2024-01-15", 150.25, 100),
    ("AAPL", "2024-01-15", 150.25, 200),
    ("AAPL", "2024-01-15", 150.10, 150),
    ("MSFT", "2024-01-15", 320.50, 80),
]

trades_df = spark.createDataFrame(
    trades_data,
    ["symbol", "date", "price", "quantity"]
)

w_rank = Window.partitionBy("symbol").orderBy(F.col("price").desc())

trades_ranked = trades_df.withColumn("row_num", F.row_number().over(w_rank)) \
    .withColumn("rank_val", F.rank().over(w_rank)) \
    .withColumn("dense_rank_val", F.dense_rank().over(w_rank))

trades_ranked.show()

Output:

For most use cases—finding the top N records per group, or deduplicating—row_number() is sufficient and slightly faster. Use rank() or dense_rank() when business rules explicitly require handling ties.

Running totals and cumulative calculations

Window functions shine when we need cumulative metrics: running balance, cumulative revenue, or rolling averages. We define a window that includes all rows from the start up to the current row.

In a ride-hailing platform, compute cumulative earnings per driver within each day:

rides_data = [
    ("driver_1", "2024-01-15", 15.50),
    ("driver_1", "2024-01-15", 22.00),
    ("driver_1", "2024-01-15", 18.75),
    ("driver_2", "2024-01-15", 25.00),
    ("driver_1", "2024-01-16", 20.00),
]

rides_df = spark.createDataFrame(
    rides_data,
    ["driver_id", "date", "fare"]
)

# Window ordered by date, unbounded preceding to current row
w_running = Window.partitionBy("driver_id", "date") \
    .orderBy(F.col("fare")) \
    .rangeBetween(Window.unboundedPreceding, Window.currentRow)

rides_with_total = rides_df.withColumn(
    "cumulative_earnings",
    F.sum("fare").over(w_running)
).orderBy("driver_id", "date", "fare")

rides_with_total.show()

The window spec uses rangeBetween(unboundedPreceding, currentRow), which means: include all rows from the partition's start up to and including this row. The result is a running total that grows as we descend through rows.

Performance note: running totals trigger a shuffle if your partitioning key differs from the data distribution. Ensure partitions are already grouped by the partition key for best performance.

Lag and lead: accessing previous and next rows

Sometimes we need to compare a row with its neighbor. lag() retrieves a value from a prior row, and lead() retrieves from a subsequent row. Both respect the window's ordering.

In a CRM system, detect when a customer's status changes:

customer_events_data = [
    ("cust_100", "2024-01-10", "active"),
    ("cust_100", "2024-01-15", "active"),
    ("cust_100", "2024-01-20", "churned"),
    ("cust_100", "2024-01-25", "churned"),
    ("cust_200", "2024-01-05", "trial"),
    ("cust_200", "2024-01-10", "active"),
]

events_df = spark.createDataFrame(
    customer_events_data,
    ["customer_id", "event_date", "status"]
)

w_lag_lead = Window.partitionBy("customer_id").orderBy("event_date")

events_with_lag = events_df.withColumn(
    "prev_status", F.lag("status").over(w_lag_lead)
).withColumn(
    "status_changed", F.col("prev_status") != F.col("status")
).fillna(False, subset=["status_changed"])

events_with_lag.show()

This reveals which rows represent a status transition. The default behavior of lag() or lead() returns null when no prior/next row exists; you can provide a default with a second argument: F.lag("status", default_value="unknown").over(w).

Lag and lead are lightweight: they don't require a full shuffle if your data is already partitioned correctly. However, be careful with offset values—requesting offset 1000 rows back requires Spark to buffer significantly more data.

Common mistakes

• Forgetting to order the window. If you don't specify orderBy(), the row order is undefined (technically the natural file order, but unreliable). Always include orderBy() when the sequence matters.
• Using window functions in a filter condition without materializing first. Filtering immediately after a window operation can cause Spark to push the filter below the window, breaking the logic. Materialize the window result or use a subquery.
• Not partitioning narrowly enough. If your partition key is too coarse (e.g., just year instead of user_id and year), you risk memory pressure as Spark tries to fit huge partitions in executor memory.
• Applying multiple overlapping windows without consolidation. If you add three separate window columns, Spark may shuffle three times. Consider building them together or in sequence on the same window definition.
• Misunderstanding rangeBetween vs rowsBetween. rangeBetween uses the ordering value (e.g., timestamps), while rowsBetween uses row count. For a running sum, rowsBetween(Window.unboundedPreceding, Window.currentRow) is often clearer.

Key takeaways

• Window functions partition and order rows without reducing row count, letting you add computed columns like rank, running total, or lag values.
• row_number(), rank(), and dense_rank() are tools for different ranking scenarios; use row_number for simple deduplication.
• Running totals and cumulative metrics rely on rangeBetween() or rowsBetween(); ensure your partition key matches your data grouping to avoid unnecessary shuffles.
• Lag() and lead() compare rows within a window; they're efficient for time-series analysis and change detection if the data is pre-partitioned.
• Always materialize window results or use subqueries before filtering; never filter immediately after a window function call.
• Monitor executor memory and shuffle volume; large partitions or many overlapping window definitions can degrade performance significantly.