Como agregar dados com groupBy no PySpark: exemplos

Agrupar linhas por uma categoria e calcular totais — vendas por região, encomendas por cliente ou valor por produto — é uma das análises mais frequentes no trabalho com dados. No PySpark esta tarefa resolve-se com o método groupBy, que funciona exatamente da mesma forma sobre mil ou mil milhões de linhas. Ao longo deste guia vais criar um DataFrame de exemplo e agregá-lo com contagens, somas e médias, terminando com uma verificação simples para garantir que os números estão certos.

Pré-requisitos

• Python 3.8 ou superior instalado.
• PySpark instalado (pip install pyspark).
• Saber ler e criar um DataFrame básico.
• Um ambiente para correr código: um notebook, o Databricks ou o terminal.

Passo 1: Criar a SparkSession e um DataFrame de exemplo

Tudo em PySpark começa com uma SparkSession, que é o ponto de entrada para trabalhar com dados distribuídos. Vamos criar uma sessão e, a partir de uma lista de tuplos, um pequeno DataFrame com vendas por região e produto. Assim temos dados concretos onde aplicar o groupBy.

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("groupby-exemplo").getOrCreate()

dados = [
    ("Norte", "Portatil", 1200),
    ("Norte", "Rato", 25),
    ("Norte", "Portatil", 1100),
    ("Sul", "Portatil", 1300),
    ("Sul", "Rato", 30),
    ("Centro", "Portatil", 1250),
]
colunas = ["regiao", "produto", "valor"]
df = spark.createDataFrame(dados, colunas)
df.show()

O createDataFrame recebe os dados e a lista de nomes de colunas. O df.show() mostra as seis linhas originais; repara que há regiões repetidas — é precisamente isso que o groupBy vai juntar num só grupo.

Passo 2: Contar linhas por grupo

A forma mais simples de agregar é contar quantas linhas existem em cada grupo. Chama groupBy com o nome da coluna e, a seguir, count().

df.groupBy("regiao").count().show()

O groupBy("regiao") devolve um objeto GroupedData — ainda não é um resultado, é apenas a definição de como agrupar. Só quando aplicas uma função como count() obténs um novo DataFrame, com uma linha por região e o número de vendas de cada uma.

Passo 3: Somar e calcular médias com agg

Contar raramente chega. Para calcular várias métricas ao mesmo tempo, usa o método agg com funções do módulo functions, que costuma importar-se como F. O alias dá um nome legível a cada coluna do resultado.

from pyspark.sql import functions as F

df.groupBy("regiao").agg(
    F.sum("valor").alias("total"),
    F.avg("valor").alias("media"),
    F.count("produto").alias("n_vendas"),
).show()

Agora cada região tem, numa só linha, o total vendido, a média por venda e o número de vendas. Sem o alias, as colunas ficariam com nomes automáticos como sum(valor), mais difíceis de reutilizar depois.

Passo 4: Agrupar por várias colunas e ordenar

Podes passar mais do que uma coluna ao groupBy para obteres um detalhe maior — por exemplo, o total por região e por produto ao mesmo tempo. Junta o orderBy para veres primeiro os maiores valores.

df.groupBy("regiao", "produto").agg(
    F.sum("valor").alias("total")
).orderBy(F.col("total").desc()).show()

O resultado tem uma linha por cada combinação de região e produto, ordenada do maior total para o menor. Trocar desc() por asc() inverte a ordem.

Verificar o resultado

Uma boa forma de confirmar que a agregação está correta é comparar a soma dos grupos com a soma de todas as linhas: as duas têm de ser iguais.

df.agg(F.sum("valor").alias("total_geral")).show()

O total geral deve dar 4905. Se somares os totais por região do Passo 3 (Norte 2325, Sul 1330 e Centro 1250) obténs exatamente 4905 — sinal de que nenhuma linha ficou de fora nem foi contada duas vezes.

Conclusão

Com groupBy, agg e o alias já consegues responder à maioria das perguntas de "totais por categoria" em PySpark, seja qual for o volume de dados. O passo seguinte natural são as window functions, que calculam agregados sem colapsar as linhas — por exemplo, um total acumulado por região. Experimenta agora: consegues adaptar o Passo 3 para mostrar também o valor máximo de cada região com F.max?