Pandas groupby — Split-Apply-Combine

TL;DR

groupby implements the split-apply-combine pattern: split the DataFrame into groups by one or more keys, apply a function to each group, combine the results back into a DataFrame or Series.

In ML, groupby is feature-engineering’s workhorse: per-user means, per-day counts, per-category z-scores, rolling aggregates within a session. Almost every “compute X for each Y” feature is a groupby.

The three operation modes:

Get the mode right and your code reads as the operation it is. Get it wrong and you’re index-aligning by hand and writing loops.

The picture in your head

A df.groupby("user_id") doesn’t actually compute anything yet — it returns a GroupBy object that knows how to iterate the rows split by key. When you call .mean() or .agg(...), that’s when the work happens, in C, group by group.

df:
    user_id  amount
    1        10
    1        20
    2        5
    2        15
    2        25
    3        100

df.groupby("user_id").sum() ->
    user_id  amount
    1        30
    2        45
    3        100

Aggregation collapses each group to one value per group.

df.groupby("user_id")["amount"].transform("mean") ->
    Series of length 6, broadcast back:
    [15.0, 15.0, 15.0, 15.0, 15.0, 100.0]
                                   # row 0,1: user 1's mean = 15
                                   # row 2,3,4: user 2's mean = 15
                                   # row 5: user 3's mean = 100

Transform keeps the original index — perfect for “add a column with the group statistic.”

agg — one row per group

The most common shape. You want a summary table.

df.groupby("user_id").agg(
    n_orders=("amount", "size"),
    total_spend=("amount", "sum"),
    avg_spend=("amount", "mean"),
    max_spend=("amount", "max"),
    last_order=("date", "max"),
)

The name=(column, function) syntax (Python 3.5+ kwargs) is the most readable form: each output column gets a name, an input column, and an aggregator. The aggregator can be a string like "mean" (calls a built-in C aggregator, fast) or any function (slower, falls back to Python).

For ML, the typical “user feature” pipeline is exactly this shape:

features = orders.groupby("user_id").agg(
    n_orders=("order_id", "nunique"),
    total_spend=("amount", "sum"),
    avg_spend=("amount", "mean"),
    days_since_last=("date", lambda d: (today - d.max()).days),
).reset_index()

This produces a per-user feature table — exactly what you’d train a model on.

transform — broadcast back to the original shape

When you want a new column that’s a function of the group:

# Z-score within each user's spending history
df["amount_z"] = df.groupby("user_id")["amount"].transform(
    lambda x: (x - x.mean()) / x.std()
)

transform returns a Series the same length as the input, with each row’s value computed from its group. The original index is preserved, so assignment back to df works.

Common patterns:

# Subtract group mean — useful for de-trending
df["amount_centered"] = df["amount"] - df.groupby("user_id")["amount"].transform("mean")

# Rank within group
df["rank"] = df.groupby("category")["score"].rank(ascending=False)

# Cumulative sum within group
df["running_total"] = df.groupby("user_id")["amount"].cumsum()

# Fraction of group total
df["share_of_user_spend"] = df["amount"] / df.groupby("user_id")["amount"].transform("sum")

Notice: cumsum, rank, cumcount are direct GroupBy methods (don’t need transform). The transform("sum") form is for built-in aggregators that return a scalar per group.

filter — keep groups by a predicate

# Only keep users with at least 5 orders
active = df.groupby("user_id").filter(lambda g: len(g) >= 5)

# Only keep categories where the median price exceeds 100
expensive = df.groupby("category").filter(lambda g: g["price"].median() > 100)

The function gets the whole group as a DataFrame; return True to keep all rows, False to drop.

For simple “filter by group size” cases, the faster pattern is:

sizes = df.groupby("user_id").size()
active_users = sizes[sizes >= 5].index
df = df[df["user_id"].isin(active_users)]

— two lines, no Python lambda per group, much faster on large data.

Multi-key groupby

df.groupby(["country", "category"])["amount"].sum()

Returns a Series with a MultiIndex. Use .reset_index() to flatten back to columns:

df.groupby(["country", "category"], as_index=False)["amount"].sum()

as_index=False is the convenient shortcut to skip the MultiIndex.

Worked example — train/test split with leak prevention

A common ML task: split a dataset by user (so a user’s records all go to the same side), and within the training set compute per-user features.

import pandas as pd
import numpy as np

orders = pd.read_parquet("orders.parquet")  # millions of rows
rng = np.random.default_rng(42)

# 1) Split users (NOT rows!) — guarantees no user appears in both sets
all_users = orders["user_id"].unique()
test_users = rng.choice(all_users, size=int(0.2 * len(all_users)), replace=False)
train_mask = ~orders["user_id"].isin(test_users)
train, test = orders[train_mask], orders[~train_mask]

# 2) Compute per-user features ONLY from training data
user_features = train.groupby("user_id").agg(
    n_orders=("order_id", "nunique"),
    avg_amount=("amount", "mean"),
    pct_returns=("is_return", "mean"),
).reset_index()

# 3) Join the features onto train and test
train = train.merge(user_features, on="user_id", how="left")
test = test.merge(user_features, on="user_id", how="left")
# Test users won't have features (they're disjoint by design)
# Fill from population stats if you need to score them
test = test.fillna({
    "n_orders": 0,
    "avg_amount": user_features["avg_amount"].median(),
    "pct_returns": 0.0,
})

Three groupbys, no leakage, no Python loops.

Performance tips

• Use built-in aggregators ("sum", "mean") over lambdas. The built-ins call optimised C code; lambdas fall back to per-group Python.
• Sort by the groupby key if you’ll groupby the same key many times. Sorted data lets Pandas use a faster algorithm.
• observed=True when grouping by category dtypes — otherwise Pandas materialises a row for every Cartesian combination of categories, even ones not in the data.
• Reach for Polars for groupby on data over a few hundred MB. Polars’s groupby is dramatically faster (and lazy by default).

Common gotchas

• NaN keys are dropped silently. Pass dropna=False to include them as a group.
• groupby returns Series for one column, DataFrame for many. df.groupby("k")["a"].sum() is a Series; df.groupby("k")[["a", "b"]].sum() is a DataFrame.
• agg with a dict is the legacy form. Prefer the named-tuple kwargs shown above; the old syntax is harder to read and clutters with MultiIndex columns.
• transform requires a return shape that matches. If you accidentally reduce, you’ll get a “Length mismatch” error.
• Ordered category dtype changes groupby ordering. Useful for pd.cut-derived buckets in a natural order.

Where this shows up in real ML codebases

• Feature engineering for tabular models. Per-user, per-item, per-session, per-day aggregates.
• Cohort analysis. Group users by sign-up week, measure retention.
• Per-class metrics. Group predictions by true label, compute precision/recall.
• Calibration. Group predictions by score bucket, compute empirical positive rate, compare to predicted.
• Data audits. Group by source / region / version, compare distributions.

The defensive habit: when you find yourself looping over unique values of a column to compute “for each X, find the…”, stop. That’s a groupby.

Resources

• Pandas user guide — Group by: split-apply-combine — pandas.pydata.org — canonical reference.
• Hadley Wickham — The Split-Apply-Combine Strategy — jstatsoft.org — the paper that named the pattern (R-flavoured but the ideas transfer).
• Polars — group by — pola.rs — faster alternative; very similar mental model.
• Modern Pandas — Tidy Data — tomaugspurger.net — when groupby is the wrong move.