Pandas18/21
Starting sandbox…
Beginner·4 graded tests

Group By (groupby)

groupby

Core Principle: The groupby method 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.
  • size vs count confusion when NaNs are present.
  • apply for everything.
  • Forgetting observed=True/False with 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 size alongside.

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.

exerciseEvery submission runs independently in a fresh namespace.

Test cases · 4

#viainputexpected stdout
1stdinsample 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
2stdinsample 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
3stdinsample gene count S1 TP53 100 TP53 rows=1 counted=1 samples=1 total=100
4stdinsample 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.