Group By (groupby)
groupby
Core Principle: The
groupbymethod implements the split-apply-combine paradigm, partitioning DataFrames by categorical keys to perform grouped aggregations, broadcasts, filters, or custom transformations.
The four apply modes
The mode determines the output shape. This is the distinction to internalise.
g = df.groupby("gene")
g["count"].agg("mean") # one row per GROUP
g["count"].transform("mean") # one row per ORIGINAL ROW (broadcast back)
g.filter(lambda x: len(x) > 5) # SUBSET of original rows
g.apply(func) # anything; slowest; pandas guesses the shape
agg
df.groupby("gene")["count"].mean()
df.groupby("gene")["count"].agg(["count", "mean", "std", "min", "max"])
df.groupby(["gene", "sample"])["count"].sum()
# named aggregation: the clearest form, and it names your output columns
df.groupby("gene").agg(
n_variants=("pos", "size"),
mean_af=("af", "mean"),
max_impact=("impact", "max"),
n_samples=("sample", "nunique"),
)
# different functions per column
df.groupby("gene").agg({"af": ["mean", "max"], "depth": "median"})
Named aggregation (new_name=("col", "func")) is the modern idiom. It produces flat, well-named columns instead of a MultiIndex you then have to flatten.
Use the string names ("mean", "sum", "size", "nunique", "first", "idxmax") rather than lambdas or np.mean. The strings dispatch to compiled Cython implementations; a lambda runs Python once per group. On 20,000 gene groups that is a 50–100× difference.
size vs count
g.size() # rows per group, INCLUDING NaN
g.count() # non-NaN values per column
g["col"].count() # non-NaN in that column
g["col"].nunique() # distinct values
size counts rows; count counts values. They differ exactly where you have missing data: which is where you most want to know.
transform
The one people underuse. It returns a result the same length as the input, broadcast back to every row of its group.
# centre expression within each gene
df["centred"] = df["count"] - df.groupby("gene")["count"].transform("mean")
# z-score within batch
df["z"] = df.groupby("batch")["value"].transform(lambda s: (s - s.mean()) / s.std())
# proportion of each sample's total
df["frac"] = df["count"] / df.groupby("sample")["count"].transform("sum")
# rank within group
df["rank_in_gene"] = df.groupby("gene")["af"].rank(ascending=False)
The alternative: aggregate, then merge back: is three lines and slower:
means = df.groupby("gene")["count"].mean().rename("gene_mean")
df = df.merge(means, on="gene")
df["centred"] = df["count"] - df["gene_mean"]
Any time you find yourself writing that pattern, transform is the answer. Per-batch normalization, per-sample scaling, within-gene centring: all transform.
filter
df.groupby("gene").filter(lambda g: len(g) > 10)
df.groupby("gene").filter(lambda g: g["af"].max() > 0.01)
Keeps whole groups based on a group-level predicate. It runs a Python lambda per group, so for simple count-based filtering the value_counts + isin version is much faster:
counts = df["gene"].value_counts()
df.loc[df["gene"].isin(counts[counts > 10].index)]
apply: last resort
df.groupby("gene").apply(lambda g: g.nlargest(3, "af")) # slow
apply is flexible and slow. Before reaching for it, check whether the operation is really:
| Want | Instead of apply |
|---|---|
| top N per group | df.sort_values(...).groupby(...).head(n) |
| best row per group | df.loc[df.groupby("g")["x"].idxmin()] |
| several aggregations | named agg |
| per-group normalization | transform |
| per-group cumulative | df.groupby("g")["x"].cumsum() |
apply also has awkward semantics: pandas inspects the return value to decide the output shape, which changes if your data changes. Pass include_groups=False in modern pandas to avoid the deprecation warning about the grouping column appearing in the passed frame.
Grouping keys
df.groupby("gene") # a column
df.groupby(["chrom", "gene"]) # several
df.groupby(level="chrom") # an index level
df.groupby(df["pos"] // 1_000_000) # a derived Series: 1 Mb bins
df.groupby(pd.cut(df["af"], bins=[0, .01, .05, 1])) # binned
df.groupby("chrom", observed=True) # categoricals: only present values
df.groupby("gene", dropna=False) # keep NaN keys as a group
df.groupby("gene", sort=False) # skip sorting: faster
df.groupby("gene", as_index=False) # keys as columns, not the index
sort=False is free speed when you do not need ordered output. as_index=False saves a reset_index().
dropna=False matters: by default, rows with a missing group key are silently dropped. If 5% of your variants have no gene annotation, they vanish from every per-gene summary without a word.
Iterating
for gene, group in df.groupby("gene"):
... # group is a DataFrame
groups = dict(list(df.groupby("gene")))
df.groupby("gene").get_group("TP53")
Fine for plotting per group or writing per-group files. Not fine as a substitute for agg/transform.
Bioinformatics examples
# variants per gene, with a breakdown
summary = variants.groupby("gene").agg(
n=("pos", "size"),
n_rare=("gnomad_af", lambda s: (s.fillna(0) < 0.001).sum()),
max_impact=("impact", "max"), # works because impact is an ORDERED category
n_samples=("sample_id", "nunique"),
).sort_values("n", ascending=False)
# mean coverage per exon
cov.groupby(["gene", "exon"])["depth"].agg(["mean", "min"])
# per-sample QC
reads.groupby("sample").agg(
total=("read_id", "size"),
mapped=("mapped", "sum"),
dup_rate=("is_dup", "mean"),
median_mapq=("mapq", "median"),
)
# normalise counts within each sample (CPM)
long["cpm"] = long["count"] / long.groupby("sample")["count"].transform("sum") * 1e6
# most significant variant per gene
best = variants.loc[variants.groupby("gene")["pvalue"].idxmin()]
# variants per 1 Mb bin, for a Manhattan-style plot
variants.groupby(["chrom", variants["pos"] // 1_000_000]).size()
The max_impact=("impact", "max") line only produces the right answer because impact is an ordered categorical: otherwise max is alphabetical and returns MODIFIER. See Categorical dtype.
Performance
df.groupby("gene", sort=False)["x"].mean() # skip sorting
df["gene"] = df["gene"].astype("category") # group on int codes
df.groupby("gene")["x"].agg("mean") # NOT .agg(lambda s: s.mean())
Rough ordering, fastest first: built-in string aggregations → transform with a string → agg with a lambda → filter → apply.
Common mistakes
- Lambdas where a string name exists. 50–100× slower.
- agg-then-merge instead of
transform. - NaN group keys silently dropped. Pass
dropna=False. sizevscountconfusion when NaNs are present.applyfor everything.- Forgetting
observed=True/Falsewith categorical keys. - Grouping on a float column. Floating-point equality makes the groups unreliable. Round or bin first.
- Not checking group sizes. A "mean" over a group of size 1 is a data point, not an estimate. Always aggregate
sizealongside.
See also
Split-Apply-Combine · Window Functions · Merging and Joining · Categorical dtype · Axes and Reductions · pandas Performance · Tidy Data
Exercise
Write summarise(df) for a long expression table with columns sample, gene
and count. Some rows have no gene annotation and some have no count. Return
one report line per gene group, formatted
<gene> rows=<n> counted=<n> samples=<n> total=<n>, where rows is the number
of rows in the group, counted is the number of rows in the group with a
non-missing count, samples is the number of distinct samples, and total is
the integer sum of the group's counts.
Rows with a missing gene must not vanish: group them under the label
unassigned. Order the lines by total descending and by gene name ascending
where totals tie. A group whose counts are all missing has a total of zero and
must still appear. Return ["no rows"] for a table with no rows at all.
Test cases · 4
| # | via | input | expected stdout |
|---|---|---|---|
| 1 | stdin | sample gene count S1 TP53 1204 S2 TP53 1150 S1 BRCA1 455 S2 BRCA1 S1 EGFR 3021 S2 900 S3 45 S3 TP53 2310 | TP53 rows=3 counted=3 samples=3 total=4664 EGFR rows=1 counted=1 samples=1 total=3021 unassigned rows=2 counted=2 samples=2 total=945 BRCA1 rows=2 counted=1 samples=2 total=455 |
| 2 | stdin | sample gene count S1 MYC S2 MYC S1 KRAS 880 | KRAS rows=1 counted=1 samples=1 total=880 MYC rows=2 counted=0 samples=2 total=0 |
| 3 | stdin | sample gene count S1 TP53 100 | TP53 rows=1 counted=1 samples=1 total=100 |
| 4 | stdin | sample gene count | no rows |
Hints · 2
01Hint
Named aggregation — agg(rows=("count", "size"), counted=("count", "count"), ...) — gives flat, well-named output columns.
02Hint
size counts rows and count counts non-missing values; they differ exactly where the data is missing.