PySpark: como remover uma coluna de um DataFrame

O PySpark é particularmente útil quando trabalhamos com grandes conjuntos de dados, pois fornece métodos eficientes para limpar nosso conjunto de dados. Neste artigo, vamos nos concentrar em uma tarefa de limpeza comum: como remover colunas de um DataFrame usando os métodos .drop() e .select() do PySpark.

Para saber mais sobre o PySpark, confira este curso de Introdução ao PySpark.

Por que remover colunas em DataFrames do PySpark?

Eliminar colunas é uma tarefa comum no pré-processamento de dados por vários motivos. A remoção de colunas desnecessárias ou irrelevantes pode reduzir consideravelmente a quantidade de dados que precisamos processar. Isso simplifica o conjunto de dados, facilitando a análise e possivelmente melhorando o desempenho de nossos trabalhos no PySpark.

Além disso, ao preparar dados para modelos de aprendizado de máquina, a seleção de variáveis independentes (features) geralmente envolve a eliminação de colunas que não contribuem para a precisão do modelo ou que podem introduzir ruído. Essa etapa pode levar a modelos mais eficientes e eficazes.

Por fim, eliminar colunas pode ser uma estratégia para melhorar a qualidade dos dados. Se uma coluna contiver uma alta porcentagem de valores ausentes ou dados errôneos, removê-la pode ser o melhor plano de ação.

Como eliminar uma única coluna de um DataFrame do PySpark

Suponha que você tenha um DataFrame df com cinco colunas: player_name, player_position, team, minutes_played e score. A coluna minutes_played tem muitos valores ausentes, por isso queremos eliminá-la.

No PySpark, podemos remover uma única coluna de um DataFrame usando o método .drop().

A sintaxe é df.drop("column_name"), onde:

• df é o DataFrame do qual queremos eliminar a coluna
• column_name é o nome da coluna a ser removida.

O método df.drop() retorna um novo DataFrame com as colunas especificadas removidas.

É assim que podemos eliminar uma coluna:

df_dropped = df.drop("minutes_played")

Como eliminar várias colunas de um DataFrame do PySpark

Vamos imaginar que queremos analisar nosso DataFrame df, mas não queremos considerar o time nem a posição do jogador. Assim, queremos eliminar essas duas colunas. Há mais de uma maneira de fazer isso, portanto vamos analisar as opções.

Para eliminar várias colunas de um DataFrame PySpark, podemos passar uma lista de nomes de colunas no método .drop(). Podemos fazer isso de duas maneiras:

# Option 1: Passing the names as a list
df_dropped = df.drop(["team", "player_position"])

# Option 2: Passing the names as separate arguments
df_dropped = df.drop("team", "player_position")

Como alternativa, podemos encadear várias chamadas .drop(). Embora tecnicamente funcional, geralmente não é considerada a abordagem mais eficiente ou elegante. Cada chamada cria um novo DataFrame, o que pode gerar sobrecarga, especialmente quando trabalhamos com conjuntos de dados maiores.

df_dropped = df.drop("team").drop("player_position")

Eliminação condicional de colunas

Em vez de especificar quais colunas devem ser eliminadas, podemos reverter o método e selecionar apenas aquelas que atendam a uma condição ou requisito. Dessa forma, o DataFrame retornado deixará de conter as colunas indesejadas. 

Usando .select()

Podemos usar o método .select() no PySpark junto com uma compreensão de lista para eliminar uma ou mais colunas com base em determinados critérios. O método é flexível, e podemos usá-lo em condições simples ou complexas. 

Esse método segue a sintaxe df.select([col for col in df.columns if condition]), em que:

• df é o DataFrame com o qual estamos trabalhando.
• condition é o critério usado para filtrar as colunas que desejamos manter.

Vamos trabalhar novamente com o DataFrame df e selecionar todas as colunas, exceto a coluna team:

df_sel = df.select([col for col in df.columns if col != "team"])

Condições complexas com .selectExpr()

Quando estamos familiarizados com o SQL e precisamos aplicar condições mais complexas ao filtrar colunas, o método .selectExpr() do PySpark oferece uma solução muito eficaz. Ele permite usar expressões do tipo SQL para selecionar e manipular colunas diretamente no código PySpark.

Veja este exemplo:

# Select specific columns and create a new 'FullMatch' column
df_sel = df.selectExpr("player_name", "player_position", "minutes_played >= 60 as FullMatch")

Aqui não estamos apenas excluindo uma coluna. Estamos:

1. Selecionando as colunas player_name e player_position.
2. Criando uma nova coluna chamada FullMatch que indica se um jogador jogou uma partida completa (60 minutos ou mais).

Essa flexibilidade torna .selectExpr() uma opção muito valiosa para situações em que a simples compreensão de listas pode não ser suficiente.

Práticas recomendadas

Ao decidir eliminar colunas do DataFrame, é fundamental seguir algumas práticas recomendadas para manter a integridade dos dados, a legibilidade do código e a eficiência.

1. Faça backup dos dados: antes de fazer qualquer modificação no DataFrame, especialmente ao eliminar colunas, é aconselhável fazer um backup. Isso garante que você possa recuperar os dados originais, se necessário.

df_backup = df.persist()  # Cache the DataFrame to avoid recomputing it later

2. Elimine com inplace=False (padrão): por padrão, o método .drop() retorna um novo DataFrame sem modificar o original. Geralmente é mais seguro. Se tiver certeza absoluta de que deseja modificar o DataFrame original, você pode usar inplace=True, mas tenha cuidado.

df_dropped = df.drop("column_to_drop", inplace=False) # Default behavior

3. Documente suas decisões: se estiver trabalhando em um projeto colaborativo ou se pretende voltar ao código no futuro, é útil incluir comentários explicando por que você optou por eliminar determinadas colunas. Essa documentação pode disponibilizar um contexto valioso para você e outras pessoas.

4. Pense na possibilidade de renomear as colunas: Em vez de eliminar colunas, devemos avaliar se renomeá-las seria mais adequado, principalmente se as informações contidas puderem ser valiosas para análises futuras. O PySpark nos fornece o método .withColumnRenamed(), que ajuda a renomear colunas.

Conclusão

Neste tutorial, você aprendeu a remover uma ou várias colunas usando os métodos .drop() e .select(). Também descrevemos métodos alternativos para utilizar expressões SQL se precisarmos de técnicas de seleção mais complexas.

Se você deseja expandir suas habilidades e fazer uso de técnicas avançadas do PySpark, confira estes cursos e materiais: 

• Big Data com PySpark
• Fundamentos de Big Data com PySpark
• Aprendizado de máquina com PySpark