Aggregations & Grouping

To aggregate is to turn many rows into one number, such as a count or a sum. groupBy does it once for each group, such as each customer.

You already know the idea from the SQL chapters (GROUP BY). This chapter shows the same work with DataFrames, on the shop orders.

groupBy moves rows between tasks, so it is a wide step with a shuffle. Spark first adds up inside each task, then sends only those partial results (chapter 5).

groupBy names the groups. agg says what to compute for each group:

from pyspark.sql import functions as F

(orders
 .groupBy("customer_id")
 .agg(F.count("*").alias("orders"), F.sum("amount").alias("total"))
 .orderBy("customer_id")
 .show())

Without orderBy, the groups can come back in any order.

For a count only, groupBy(...).count() is a short form. groupBy("status").count() gives paid 4, pending 1 and refunded 1.

Filter the groups

To keep only some groups, filter after agg. This works like HAVING in SQL:

(orders
 .groupBy("customer_id")
 .agg(F.sum("amount").alias("total"))
 .filter(F.col("total") > 50)
 .orderBy("customer_id")
 .show())

One agg can compute many numbers in a single pass over the data. Here are six for each status:

(orders
 .groupBy("status")
 .agg(
     F.count("*").alias("orders"),
     F.count("amount").alias("with_amount"),
     F.sum("amount").alias("total"),
     F.round(F.avg("amount"), 2).alias("avg"),
     F.max("amount").alias("biggest"),
     F.countDistinct("customer_id").alias("customers"),
 )
 .orderBy("status")
 .show())

Look at the paid row. avg is 160.5 / 4 = 40.125, and round(..., 2) gives 40.13, because Spark rounds a half up. countDistinct counts different customers: Asha, Carla and Dev.

when inside sum gives several totals side by side, also in one pass:

orders.agg(
    F.sum(F.when(F.col("status") == "paid", F.col("amount"))).alias("paid"),
    F.sum(F.when(F.col("status") == "refunded", F.col("amount"))).alias("refunded"),
).show()

rollup adds a total row to the groups. Here the orders are joined to the customers, and grouped by country:

joined = orders.join(customers, "customer_id")

(joined
 .rollup("country")
 .agg(F.sum("amount").alias("total"), F.grouping("country").alias("is_total"))
 .orderBy("is_total", "country")
 .show())

grouping("country") tells them apart: it is 1 on the total row and 0 on a real group. Without it, a report cannot show which NULL is the total.

With two columns, rollup("country", "status") adds a subtotal for each country and one grand total: 10 rows here. cube adds every mix of subtotals: 13 rows. The SQL chapter Sets, Pivots & Reshaping Results explains the counts.

pivot turns the values of one column into column names. Here each status becomes a column:

(orders
 .groupBy("customer_id")
 .pivot("status", ["paid", "pending", "refunded"])
 .agg(F.sum("amount"))
 .orderBy("customer_id")
 .show())

Give the list of values, as here. Without the list, Spark first runs an extra job to find all the different statuses.

To show 0 instead of NULL, add .fillna(0) after the pivot. But then an unknown amount and no order both look like 0.

Five tools turn rows into numbers:

Check yourself. Why is the average of the paid orders 40.13 and not 40.12? (Answer: 40.125 is a half, and Spark rounds a half up.)