Agrupar e agregar

Separar, aplicar, combinar

O groupby trabalha em três etapas: separa os dados em grupos, aplica um cálculo a cada grupo e combina os resultados em uma tabela:

import pandas as pd

pd.set_option("display.width", 170)
pd.set_option("display.max_columns", 20)

CANONICO = {"engenharia": "Engenharia", "marketing": "Marketing", "finanças": "Finanças",
            "ciência de dados": "Ciência de Dados", "rh": "RH"}


def carregar_limpo() -> pd.DataFrame:
    df = pd.read_csv("dados/funcionarios_15.csv")
    df["salario"] = pd.to_numeric(df["salario"].str.replace("R$ ", "", regex=False))
    df["departamento"] = df["departamento"].str.strip().str.lower().map(CANONICO)
    return df.drop_duplicates(subset=["id_funcionario"], keep="first").reset_index(drop=True)


df = carregar_limpo()
print(df.groupby("departamento")["salario"].mean().round(1))

departamento
Ciência de Dados    6450.0
Engenharia          5100.0
Finanças            5925.0
Marketing           5362.5
RH                  5100.0
Name: salario, dtype: float64

`size` contra `count`

As funções de agregação mais usadas são sum, mean, median, min, max, count, size e std. Duas delas parecem iguais até aparecerem ausentes: o size conta todas as linhas do grupo, e o count conta só as não ausentes:

print(df.groupby("departamento")["salario"].agg(["size", "count"]))

                  size  count
departamento                 
Ciência de Dados     2      1
Engenharia           3      3
Finanças             2      2
Marketing            4      4
RH                   2      1

Nos departamentos em que alguém não tem salário, o size é maior que o count. Qual usar depende da pergunta: "quantas pessoas" é size, "quantos salários conhecidos" é count.

A armadilha: o grupo que some

Por padrão, o groupby descarta as linhas em que a chave é ausente. Aqui, uma pessoa não tem departamento, e ela desaparece de qualquer resumo por departamento, sem aviso:

print(len(df), int(df.groupby("departamento")["id_funcionario"].size().sum()))
com_ausente = df.groupby("departamento", dropna=False)["id_funcionario"].size()
print(int(com_ausente.sum()), com_ausente.index.isna().sum())

14 13
14 1

O total de funcionários é 14, e a soma dos grupos é só 13. Com dropna=False, o grupo dos "sem departamento" aparece, e a soma fecha. Em um relatório, um funcionário faltando na contagem é um erro que ninguém percebe, e por isso eu confiro que a soma dos grupos bate com o total.

Várias colunas, várias funções

print(df.groupby(["departamento", "cidade"])["salario"].mean().round(0).head(5))
print(df.groupby("departamento")[["salario", "nota_desempenho"]].mean().round(2))

departamento      cidade        
Ciência de Dados  Rio de Janeiro    6450.0
Engenharia        Porto Alegre      4800.0
                  São Paulo         5250.0
Finanças          Porto Alegre      5925.0
Marketing         Belo Horizonte    5025.0
Name: salario, dtype: float64
                  salario  nota_desempenho
departamento                              
Ciência de Dados   6450.0             6.00
Engenharia         5100.0             6.00
Finanças           5925.0             4.50
Marketing          5362.5             6.75
RH                 5100.0             7.50

O primeiro agrupa por duas chaves (um grupo para cada combinação), e o segundo calcula a média de duas colunas ao mesmo tempo. O agg calcula várias funções de uma vez, e a agregação nomeada deixa nomes legíveis nas colunas do resultado, em vez dos cabeçalhos aninhados:

resumo = df.groupby("departamento").agg(
    salario_medio=("salario", "mean"),
    maior_salario=("salario", "max"),
    nota_media=("nota_desempenho", "mean"),
    funcionarios=("id_funcionario", "count"),
)
print(resumo.round(1))

                  salario_medio  maior_salario  nota_media  funcionarios
departamento                                                            
Ciência de Dados         6450.0         6450.0         6.0             2
Engenharia               5100.0         6000.0         6.0             3
Finanças                 5925.0         6300.0         4.5             2
Marketing                5362.5         6150.0         6.8             4
RH                       5100.0         5100.0         7.5             2

O grupo virou índice

Depois de agrupar, a coluna do grupo vira o índice do resultado, e não uma coluna comum. Acessá-la como coluna dá um KeyError. O reset_index() a devolve (ou, desde o início, o as_index=False):

try:
    resumo["departamento"]
except KeyError as erro:
    print("KeyError:", erro)

print(resumo.reset_index().columns.tolist())
print(df.groupby("departamento", as_index=False)["salario"].mean().columns.tolist())

KeyError: 'departamento'
['departamento', 'salario_medio', 'maior_salario', 'nota_media', 'funcionarios']
['departamento', 'salario']

Use o reset_index() no momento em que você precisar tratar o resultado como uma tabela comum: ordenar por ele, filtrar ou juntar com outra tabela.

Um cálculo por grupo, mantendo a forma original

O agg reduz cada grupo a uma linha. O transform calcula por grupo, mas devolve uma linha para cada linha original, o que permite comparar cada pessoa com a média do seu departamento:

media_do_depto = df.groupby("departamento")["salario"].transform("mean")
df["acima_da_media_do_depto"] = df["salario"] > media_do_depto
print(df[["nome", "departamento", "salario", "acima_da_media_do_depto"]].head(5))

              nome      departamento  salario  acima_da_media_do_depto
0        Ana Souza        Engenharia   4500.0                    False
1       Bruno Lima         Marketing   4650.0                    False
2     Carla Mendes        Engenharia   4800.0                    False
3      Diego Rocha  Ciência de Dados      NaN                    False
4  Elisa Fernandes                RH   5100.0                    False

O capítulo 23 aprofunda o transform, o filter e o apply por grupo.

O groupby escondido em toda pergunta de negócio

"Qual região vende mais?", "Qual categoria tem a maior margem?", "Qual segmento cancela mais?": quase toda pergunta de negócio é um groupby disfarçado. Quando uma pergunta tem a forma "o quê, por quê", é um agrupamento.