4. Grouping & aggregation

Intermediate · 10 min read

groupby answers questions like "average latency per model" — split the rows into groups, apply a calculation, combine the results.

import pandas as pd

df = pd.DataFrame({
    "model":      ["gpt-4o-mini", "llama-3.1-8b", "gpt-4o", "gpt-4o-mini", "llama-3.1-8b", "gpt-4o"],
    "question":   ["What is RAG?", "What is RAG?", "What is RAG?", "Explain agents", "Explain agents", "Explain agents"],
    "latency_ms": [820, 410, 1350, 960, 450, 1500],
    "tokens":     [210, 190, 240, 260, 230, 300],
    "cost_usd":   [0.00013, 0.00002, 0.0024, 0.00016, 0.00002, 0.003],
    "rating":     [5, 4, 5, 4, 3, 5],
})

4.1 One statistic per group

print(df.groupby("model")["latency_ms"].mean())
print(df.groupby("model")["rating"].agg(["mean", "min", "max"]))

Output

model
gpt-4o          1425.0
gpt-4o-mini      890.0
llama-3.1-8b     430.0
Name: latency_ms, dtype: float64
              mean  min  max
model                       
gpt-4o         5.0    5    5
gpt-4o-mini    4.5    4    5
llama-3.1-8b   3.5    3    4

4.2 Named aggregations — a clean report

report = (
    df.groupby("model")
      .agg(answers=("question", "count"),
           avg_rating=("rating", "mean"),
           p50_latency=("latency_ms", "median"),
           total_tokens=("tokens", "sum"),
           total_cost=("cost_usd", "sum"))
      .sort_values("avg_rating", ascending=False)
)
print(report.round(5))

Output

              answers  avg_rating  p50_latency  total_tokens  total_cost
model                                                                   
gpt-4o              2         5.0       1425.0           540     0.00540
gpt-4o-mini         2         4.5        890.0           470     0.00029
llama-3.1-8b        2         3.5        430.0           420     0.00004

4.3 Group by several columns

print(df.groupby(["question", "model"])["rating"].mean().unstack())   # unstack: models → columns

Output

model           gpt-4o  gpt-4o-mini  llama-3.1-8b
question                                         
Explain agents     5.0          4.0           3.0
What is RAG?       5.0          5.0           4.0

4.4 transform — a group value on every row

agg returns one row per group; transform returns a value for each original row — perfect for "compared with its group" columns.

df["model_avg_latency"] = df.groupby("model")["latency_ms"].transform("mean")
df["vs_model_avg"] = df["latency_ms"] - df["model_avg_latency"]
df["rank_in_question"] = df.groupby("question")["rating"].rank(ascending=False, method="min")
print(df[["model", "question", "latency_ms", "vs_model_avg", "rank_in_question"]])

Output

          model        question  latency_ms  vs_model_avg  rank_in_question
0   gpt-4o-mini    What is RAG?         820         -70.0               1.0
1  llama-3.1-8b    What is RAG?         410         -20.0               3.0
2        gpt-4o    What is RAG?        1350         -75.0               1.0
3   gpt-4o-mini  Explain agents         960          70.0               2.0
4  llama-3.1-8b  Explain agents         450          20.0               3.0
5        gpt-4o  Explain agents        1500          75.0               1.0

4.5 Counting

print(df["model"].value_counts())
print(df["rating"].value_counts(normalize=True).sort_index())   # shares instead of counts
print(df.groupby("model").size())

Output

model
gpt-4o-mini     2
llama-3.1-8b    2
gpt-4o          2
Name: count, dtype: int64
rating
3    0.166667
4    0.333333
5    0.500000
Name: proportion, dtype: float64
model
gpt-4o          2
gpt-4o-mini     2
llama-3.1-8b    2
dtype: int64

4.6 Pivot tables and crosstab

pivot = pd.pivot_table(df, index="model", columns="question", values="latency_ms", aggfunc="mean")
print(pivot)
print(pd.crosstab(df["model"], df["rating"]))            # counts of each rating per model

Output

question      Explain agents  What is RAG?
model                                     
gpt-4o                1500.0        1350.0
gpt-4o-mini            960.0         820.0
llama-3.1-8b           450.0         410.0
rating        3  4  5
model                
gpt-4o        0  0  2
gpt-4o-mini   0  1  1
llama-3.1-8b  1  1  0

4.7 Filter whole groups

reliable = df.groupby("model").filter(lambda g: g["rating"].min() >= 4)   # keep models never rated below 4
print(reliable["model"].unique().tolist())

Output

['gpt-4o-mini', 'gpt-4o']

Practice

• For each question, which model answered fastest? (Hint: idxmin on a grouped column.)
• Add a column with each answer's share of its model's total tokens.

Next: Combining & reshaping →