13. Conditional aggregation and `pivot`

Two techniques for turning long data into wide summaries.

Conditional aggregation with when

Sometimes you want to sum only some rows within each group — revenue from one category, count of a certain status. Wrap the value in F.when(...) so non-matching rows contribute nothing:

orders.dropDuplicates().groupBy("country").agg(
    F.sum(F.when(F.col("category") == "electronics", F.col("quantity")).otherwise(0)).alias("electronics_qty"),
    F.sum(F.when(F.col("category") == "grocery", F.col("quantity")).otherwise(0)).alias("grocery_qty"),
).show()

when(condition, value).otherwise(other) is PySpark's CASE WHEN. Here each sum adds the quantity only for its category and 0 elsewhere — so one pass over the data produces several category columns side by side.

pivot — categories become columns

pivot does the same reshape declaratively: pick a column whose values become columns. Count orders per country, broken out by category:

orders.dropDuplicates().groupBy("country").pivot("category").agg(F.count("*")).show()

Each distinct category value (books, electronics, …) becomes its own column; empty cells are null (no orders of that category from that country). Pivots are perfect for cross-tabs and reports — a country × category grid in one line.

Tip. If you know the values you want as columns, pass them: pivot("category", ["electronics", "grocery"]). It's faster (Spark skips a scan to discover them) and pins the column order.

Your turn

Build a country × category grid of total quantity: group deduplicated orders by country, pivot on category, and aggregate sum("quantity"). Assign to result.

result = orders  # <- dropDuplicates, groupBy country, pivot category, agg sum(quantity)
result.show()

result = (
    orders.dropDuplicates()
    .groupBy("country")
    .pivot("category")
    .agg(F.sum("quantity"))
)