pandas
groupby
python
data analysis
lambda functions

Group by with multiple columns using lambda

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Grouping by multiple columns in pandas is straightforward: pass a list of grouping keys to groupby. The lambda part comes later, when you need a custom aggregation, transformation, or filter over each group. The important thing is to choose the right post-group operation, because agg, transform, and apply do different jobs.

Basic Grouping by Multiple Columns

Suppose you want to group sales by region and product.

python
1import pandas as pd
2
3df = pd.DataFrame(
4    {
5        "region": ["East", "East", "West", "West", "East"],
6        "product": ["A", "B", "A", "B", "A"],
7        "sales": [100, 120, 90, 150, 130],
8        "quantity": [2, 3, 1, 4, 2]
9    }
10)
11
12grouped = df.groupby(["region", "product"])["sales"].sum()
13print(grouped)

The list ['region', 'product'] creates a multi-key grouping. Each unique pair becomes one group.

Use a Lambda with agg

If you need a custom summary, attach a lambda through agg.

python
1summary = df.groupby(["region", "product"]).agg(
2    sales_total=("sales", "sum"),
3    sales_range=("sales", lambda s: s.max() - s.min()),
4    avg_quantity=("quantity", lambda q: q.mean())
5)
6
7print(summary)

This is a good pattern when each group should collapse into one output row with one or more summary values.

Use a Lambda with transform

transform is different. It returns a result aligned to the original rows, which makes it useful when you want a per-row value derived from each group.

python
1import pandas as pd
2
3df = pd.DataFrame(
4    {
5        "region": ["East", "East", "West", "West", "East"],
6        "product": ["A", "B", "A", "B", "A"],
7        "sales": [100, 120, 90, 150, 130]
8    }
9)
10
11df["group_share"] = (
12    df["sales"] /
13    df.groupby(["region", "product"])["sales"].transform(lambda s: s.sum())
14)
15
16print(df)

Here, each row gets its share of the total sales for its own (region, product) group.

Use a Lambda with filter

If you want to keep only groups that satisfy a condition, filter is the appropriate tool.

python
1filtered = df.groupby(["region", "product"]).filter(
2    lambda g: g["sales"].sum() >= 200
3)
4
5print(filtered)

This preserves the original rows, but only for groups that pass the lambda condition.

When a Named Function Is Better

Lambdas are concise, but they are not always the clearest option. If the group logic is more than a short expression, write a named function instead.

python
1def coefficient_of_variation(series):
2    mean = series.mean()
3    if mean == 0:
4        return 0.0
5    return series.std(ddof=0) / mean
6
7result = df.groupby(["region", "product"]).agg(
8    sales_cv=("sales", coefficient_of_variation)
9)

This is easier to test and easier to read than a dense lambda with several intermediate steps.

Common Pitfalls

A common mistake is using apply for everything. apply is flexible, but it is often slower and harder to reason about than agg, transform, or filter. Use the most specific group operation that matches the job.

Another mistake is forgetting what shape the result should have. agg reduces groups, transform preserves row count, and filter keeps or removes whole groups. If you choose the wrong tool, the output shape will look wrong even if the code runs.

Developers also sometimes write lambdas that depend on external mutable state. That makes group operations harder to debug and can lead to surprising results. Keep the lambda focused on the group data passed into it.

Finally, remember that grouping by multiple columns produces a MultiIndex by default in many cases. If you want plain columns afterward, call reset_index().

Summary

  • Group by multiple columns in pandas by passing a list of column names to groupby.
  • Use agg for one-row-per-group summaries, often with small lambda functions.
  • Use transform when the result must align back to the original rows.
  • Use filter to keep or discard entire groups based on a condition.
  • Prefer a named function over a lambda when the group logic becomes complex.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.