Filtragem e Seleção de Dados com Python

Na maioria dos cenários do mundo real, não é necessário trabalhar com todas as colunas ou linhas de um conjunto de dados. Tabelas grandes podem ter centenas de colunas e milhões de linhas, sendo que a maioria pode ser irrelevante para a análise específica. Neste capítulo, será apresentado como "extrair" exatamente os dados necessários utilizando os métodos select() e filter().

Seleção de Colunas Específicas

O método select() permite criar um novo DataFrame que contém apenas as colunas escolhidas. Isso reduz a quantidade de memória utilizada pelo cluster e torna os resultados muito mais fáceis de interpretar.

# Select only two specific columns
selected_df = df.select("carat", "price")

display(selected_df)

Observe que o Spark não altera o df original. Em vez disso, ele cria um novo chamado selected_df. Se você quiser selecionar todas as colunas exceto uma, ou realizar um cálculo durante a seleção, será necessário utilizar uma sintaxe mais avançada, mas para tarefas básicas, passar os nomes das colunas como strings é o método padrão.

Filtrando linhas com condições

O método filter() (ou seu alias where()) funciona como um filtro para seus dados. Você fornece uma condição, e o Spark mantém apenas as linhas em que essa condição é verdadeira.

# Filter for rows where the carat is '0.23'
filtered_df = df.filter(df.carat == 0.23)

display(filtered_df)

Você pode usar operadores de comparação padrão como == (igual), != (diferente), > (maior que) e < (menor que).

Combinando Múltiplos Filtros

Frequentemente, é necessário aplicar mais de uma regra ao mesmo tempo. Para isso, é possível encadear filtros ou usar operadores lógicos como "E" (&) e "Ou" (|).

good_carat_df = df.filter((df.carat == 0.23) & (df.cut == "Good"))

display(good_carat_df)

Seleção e filtragem em uma única etapa

Como o Spark utiliza uma API "fluent", é possível encadear esses comandos em uma única linha de código. Esse é um padrão muito comum em engenharia de dados profissional:

# Select specific diamonds by `price > 500`
high_profit_df = df.select("carat", "price").filter(df.price > 500)

display(high_profit_df)

Verificação do Trabalho

Após cada seleção ou filtro, é uma boa prática executar um count(). Se você começa com 10.000 linhas e após um filtro tem 0, isso indica que sua condição de filtro pode estar muito restritiva ou conter um erro de digitação nos valores de string.