One Row Per Group

The Daily Grind's owner doesn't want to read eleven orders. She wants three lines: how each size sold. Many rows in, one row per size out. That is aggregation, and in Spark it is always the same two moves: say what the groups are, then say what to compute for each.

Watch it happen

The orders table, eleven rows, five of them large:

from pyspark.sql import functions as F

result = orders.groupBy("size").agg(
    F.sum("price").alias("revenue"),
    F.count("*").alias("orders"),
    F.round(F.avg("tip"), 2).alias("avg_tip"),
)

groupBy("size") decides the rows of the answer: one per distinct size. agg(...) decides the columns: each aggregate looks at every order in the group and returns one value. Eleven rows became three, and nothing about an individual order survives except through a total.

Name every aggregate

Leave off .alias and the engine names the columns for you: sum(price), count(1), avg(tip). The grader matches columns by name, so an unnamed aggregate is a wrong answer however right the numbers are. Name them.

count("*") is not count("col")

count("*") counts rows. count("tip") counts rows where tip is not null. On data with gaps the two differ, and the second is usually the one a question means by "how many have a tip". sum and avg skip nulls too: an average of 1.0 and null is 1.0, not 0.5.

The other reducers

sum, count and avg are three of a family. F.max and F.min return the largest and smallest value in the group, and F.countDistinct("customer_id") counts how many different values there are: three orders from the same regular are three rows and one buyer. Each looks at the group the same way, and each needs its .alias:

orders.groupBy("size").agg(
    F.max("price").alias("priciest"),
    F.min("price").alias("cheapest"),
    F.countDistinct("drink").alias("drinks_on_offer"),
)

Counting only some of the rows

"How many orders in each size were tipped?" is a count with a condition inside it. The pattern is a when that yields 1 or 0, summed:

orders.groupBy("size").agg(
    F.sum(F.when(F.col("tip") > 0, 1).otherwise(0)).alias("tipped"),
)

The .otherwise(0) is doing work. Leave it off and a size with no tipped orders sums nulls, which gives null, not 0. And a filter before the groupBy is not the same thing: a group with nothing left after the filter is not a 0 row, it is no row at all.

Lists, and their order

collect_list gathers a group's values into an array; collect_set gathers the distinct ones. The order of what comes back is not defined. Run collect_set("drink") three times on the same eleven rows here and you get three different orders. If the order matters to anyone, including a grader, wrap it: F.sort_array(F.collect_set("drink")).

first and last have the same problem wearing a different name. They return whichever value the engine happened to meet first, and sorting the DataFrame beforehand does not change that: orders.orderBy(F.col("price").desc()).groupBy("size").agg(F.first("drink")) returned the same drink as the unsorted version. To get "the drink with the highest price in each size" you need the window functions of a later lesson, not first.

Keep only some groups

A filter after the aggregation keeps groups, not rows. SQL calls this HAVING; here it is just filter on the aliased column:

orders.groupBy("size").agg(F.sum("price").alias("revenue")).filter(F.col("revenue") > 15)

Two groups survive: large and small. Medium, at 10.25, does not.

Nine problems, in the order they build:

1. Daily Ride Count — one group, one count.
2. Revenue Per Store — a sum per group.
3. Distinct Buyers — count distinct values, not rows.
4. Priciest Per Category — a maximum per group.
5. Barista Averages — an average, then keep only some groups.
6. Store-Day Revenue — two grouping keys.
7. Plays Per Month — group by a column you derived.
8. Count the Refunds — count with a condition, and keep the zero groups.
9. Basket Contents — a set per group, sorted so it can be graded.