Funções de janela (Window) no PySpark: passo a passo

As funções de janela (Window) no PySpark permitem numerar, ordenar e comparar linhas dentro de grupos sem perder o detalhe de cada registo. Ao contrário do groupBy, que junta tudo numa só linha por grupo, uma função de janela mantém todas as linhas e acrescenta o cálculo numa coluna nova. É assim que se criam rankings, totais acumulados e comparações com a linha anterior. Imagine que quer saber os melhores vendedores de cada região ou um total de vendas acumulado: são casos perfeitos para funções de janela.

Pré-requisitos

• Ter o PySpark instalado, ou um ambiente como Databricks ou Microsoft Fabric já configurado.
• Saber criar um DataFrame e usar withColumn.
• Conhecimentos básicos de Python.

Passo 1: Criar dados de exemplo

Começamos por criar uma SparkSession e um DataFrame pequeno, com vendas por pessoa e por região. Um exemplo pequeno ajuda a ver claramente o efeito de cada função. Guarde este código, porque vamos construir sobre ele nos passos seguintes.

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("janelas").getOrCreate()

dados = [
    ("Ana", "Norte", 100),
    ("Bruno", "Norte", 150),
    ("Carla", "Norte", 120),
    ("Diogo", "Sul", 200),
    ("Eva", "Sul", 180),
]
df = spark.createDataFrame(dados, ["nome", "regiao", "vendas"])
df.show()

Passo 2: Definir a janela

Uma janela define como o Spark deve olhar para as linhas antes de calcular. O partitionBy divide os dados em grupos — aqui, uma janela por região — e o orderBy define a ordem dentro de cada grupo, neste caso das vendas mais altas para as mais baixas. Ainda não se calcula nada: apenas descrevemos a janela que vamos reutilizar.

Esta separação entre definir a janela e aplicá-la torna o código mais legível: podemos reutilizar a mesma janela em vários cálculos sem a repetir.

from pyspark.sql.window import Window
from pyspark.sql import functions as F

janela = Window.partitionBy("regiao").orderBy(F.col("vendas").desc())

Passo 3: Numerar linhas com row_number()

A função row_number() atribui uma posição sequencial — 1, 2, 3... — dentro de cada região, seguindo a ordem da janela. É a forma mais direta de criar um ranking. Usamos withColumn para guardar o resultado numa coluna nova chamada posicao.

df = df.withColumn("posicao", F.row_number().over(janela))
df.show()

Passo 4: Calcular um total acumulado

Quando aplicamos sum() a uma janela ordenada, o Spark não soma o grupo todo de uma vez: soma de forma acumulada, linha a linha, até à linha atual. É exatamente o que precisamos para um total corrente dentro de cada região.

Por omissão, uma janela ordenada usa a moldura desde o início do grupo até à linha atual, e é isso que produz o efeito acumulado que vemos na coluna.

df = df.withColumn("acumulado", F.sum("vendas").over(janela))
df.show()

Passo 5: Comparar com a linha anterior usando lag()

A função lag() vai buscar o valor de uma linha anterior dentro da mesma janela — por omissão, a linha imediatamente acima. É muito útil para calcular variações, como a diferença de vendas entre duas posições. A primeira linha de cada grupo não tem antecessor, por isso recebe null.

df = df.withColumn("venda_anterior", F.lag("vendas").over(janela))
df.show()

Verificar o resultado

Repare no resultado final. Na região Norte, o Bruno aparece em primeiro (150), seguido da Carla (120) e da Ana (100), por isso a coluna posicao mostra 1, 2 e 3. A coluna acumulado cresce ao longo do grupo — 150, depois 270, depois 370 — e reinicia na região Sul. Já a coluna venda_anterior fica a null na primeira linha de cada região, porque não existe nenhuma linha antes dela. Se vir estes valores, as funções de janela estão a funcionar corretamente.

Para experimentar, cole cada bloco de código por ordem num notebook e execute df.show() depois de cada passo, observando as colunas novas a aparecer uma a uma.

Conclusão

Com poucas linhas de código, criou um ranking, um total acumulado e uma comparação entre linhas, tudo sem perder o detalhe dos dados originais. A partir daqui, experimente rank() e dense_rank() para lidar com empates, ou lead() para olhar para a linha seguinte. Para médias móveis, ajuste a moldura da janela com rowsBetween. Qual é o cálculo do seu dia a dia que ficaria mais simples com uma função de janela?