Python · easy

Group a DataFrame by category and compute multiple aggregations at once.

Asked in Analytics Engineer interviews, in the Python round.

Short answer

df.groupby('cat').agg({'x':'sum','y':'mean'}).

Sample data

These tables are loaded into an in-browser engine on this page, so you can run your own answer against them. No account needed.

sales

categoryregionamount
electronicseast120
electronicswest200
electronicseast90
homeeast60
homewest150
homewest80
traveleast40
travelwest110

How to answer it

Definition

One pass over the groups producing several aggregates. .agg() takes a dict or, better, named keyword arguments that also set the output column names.

Syntax

df.groupby("cat").agg(
    total=("amount", "sum"),
    average=("amount", "mean"),
    orders=("order_id", "count"),
    customers=("user_id", "nunique"),
).reset_index()

Example

summary = (
    orders
    .groupby("category", dropna=False)
    .agg(
        revenue=("amount", "sum"),
        avg_order=("amount", "mean"),
        orders=("order_id", "count"),
        customers=("user_id", "nunique"),
    )
    .reset_index()
    .sort_values("revenue", ascending=False)
)

Output

  category  revenue  avg_order  orders  customers
0     Home  91200.0      48.25    1890       1204
1  Outdoor  58800.0      52.10    1129        842

Named aggregation over a dict

Passing a dict works, but agg({"amount": ["sum", "mean"]}) produces a MultiIndex on the columns that you then flatten by hand. Named aggregation gives flat, readable column names in the same call, which is why it is the version worth memorising.

Key points

FAQ

The 30-second answer

groupby on the category, then agg with named aggregations, so each keyword is the output column name and the value is a tuple of the source column and the function.
I prefer the named form over passing a dict, because the dict version gives you a MultiIndex on the columns that you then have to flatten. I would also mention two details that bite people: count counts non-null values while size counts rows, so they disagree exactly where the nulls are, and dropna set to False if a missing category should be its own group rather than quietly disappearing from the report.

In the interview

Practice this for real