24. The window spec — `partitionBy` and `orderBy`

📖 Reading · 4 min
💡 Every code box below is live — edit it and hit Run.

Window functions compute a value across a set of related rows while keeping every row in the output. That last part is the difference from groupBy: grouping collapses rows; a window function annotates them.

Compare each product's price to its category's total — without losing a single row:

from pyspark.sql.window import Window

w = Window.partitionBy("category")

orders.dropDuplicates().withColumn(
    "category_total", F.round(F.sum("unit_price").over(w), 2)
).select("product", "category", "unit_price", "category_total").show(8)

Every row survives, and each carries its category's total alongside its own price. With groupBy("category") you'd get one row per category and lose product entirely.

The two parts of a window spec

  • partitionBy(cols) — split rows into independent groups. The function restarts for each partition. (Omit it and the whole DataFrame is one window.)
  • orderBy(cols) — order rows within each partition. Required for anything position-based (ranking, lag, running totals); optional for whole-partition aggregates like the sum above.

You attach a window to a function with .over(w). Any aggregate (sum, avg, count, max) works as a window function, plus the dedicated ones (ranking, lag/lead) in the next lessons.

groupBy vs window — pick by output shape

  • Want one row per group (a summary)? groupBy.
  • Want every original row, plus a group-level value? Window function.

A frequent pattern: window to compute each row's share of its group.

w = Window.partitionBy("category")
orders.dropDuplicates().withColumn(
    "pct_of_category",
    F.round(F.col("unit_price") / F.sum("unit_price").over(w) * 100, 1),
).select("product", "category", "unit_price", "pct_of_category").show(8)

Your turn

Add a category_avg column — the average unit_price of each row's category (rounded to 2) — to deduplicated orders, keeping all rows. Return product, category, unit_price, category_avg. Assign to result.

w = Window.partitionBy("category")
result = orders  # <- dropDuplicates, withColumn category_avg = avg(unit_price).over(w), select 4 cols
result.show()
w = Window.partitionBy("category")
result = orders.dropDuplicates().withColumn(
    "category_avg", F.round(F.avg("unit_price").over(w), 2)
).select("product", "category", "unit_price", "category_avg")