PySpark: Cómo eliminar una columna de un DataFrame

PySpark es especialmente útil cuando trabajamos con grandes conjuntos de datos, porque proporciona métodos eficientes para limpiar nuestro conjunto de datos. En este artículo, nos centraremos en una tarea de limpieza habitual: cómo eliminar columnas de un DataFrame utilizando los métodos de PySpark .drop() y .select().

Para saber más sobre PySpark, consulta el curso Introducción a PySpark.

¿Por qué eliminar columnas en los DataFrames de PySpark?

Eliminar columnas es una tarea habitual en el preprocesamiento de datos por diversas razones. Eliminar columnas innecesarias o irrelevantes puede reducir significativamente la cantidad de datos que debemos procesar. Esto optimiza nuestro conjunto de datos, facilita su análisis y puede aumentar el rendimiento de nuestros trabajos PySpark.

Además, al preparar los datos para los modelos de machine learning, la selección de características suele implicar la eliminación de columnas que no contribuyen a la exactitud del modelo o que podrían introducir ruido. Este paso puede conducir a modelos más eficientes y eficaces.

Por último, eliminar columnas puede ser una estrategia para mejorar la calidad de los datos. Si una columna contiene un alto porcentaje de valores que faltan o datos erróneos, eliminarla puede ser la mejor opción.

Cómo eliminar una sola columna de un DataFrame PySpark

Supongamos que tenemos un DataFrame df con cinco columnas: player_name, player_position, team, minutes_played y score. La columna minutes_played tiene muchos valores que faltan, por lo que queremos eliminarla.

En PySpark, podemos eliminar una sola columna de un DataFrame utilizando el método .drop().

La sintaxis es df.drop("column_name"), donde:

• df es el DataFrame del que queremos eliminar la columna
• column_name es el nombre de la columna que hay que eliminar

El método df.drop() devuelve un nuevo DataFrame con las columnas especificadas eliminadas.

Así es como podemos eliminar una columna:

df_dropped = df.drop("minutes_played")

Cómo eliminar varias columnas de un DataFrame PySpark

Imaginemos que queremos analizar nuestro DataFrame df, pero no queremos tener en cuenta el equipo ni la posición del jugador, por lo que queremos eliminar ambas columnas. Hay más de una forma de hacerlo, así que vamos a explorar nuestras opciones.

Para eliminar varias columnas de un DataFrame PySpark, podemos pasar una lista de nombres de columnas al método .drop(). Podemos hacerlo de dos maneras:

# 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")

También podemos encadenar varias llamadas a .drop(). Aunque este enfoque es técnicamente funcional, no se suele considerar el más eficiente o elegante. Cada llamada crea un nuevo DataFrame, lo que puede introducir sobrecarga, especialmente cuando se trabaja con conjuntos de datos grandes.

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

Eliminar columnas condicionalmente

En lugar de especificar qué columnas eliminar, podemos invertir el método y seleccionar solo las que cumplan una condición o un requisito. De este modo, nuestro DataFrame devuelto ya no contendrá las columnas no deseadas. 

Utilizar .select()

Podemos utilizar el método .select() de PySpark junto con una comprensión de listas para eliminar una o más columnas en función de determinados criterios. El método es flexible, y podemos utilizarlo con una condición sencilla o compleja. 

Este método sigue la sintaxis df.select([col for col in df.columns if condition]), donde:

• df es el DataFrame con el que estás trabajando.
• condition es el criterio utilizado para filtrar las columnas que quieres conservar.

Trabajemos de nuevo con nuestro DataFrame df y seleccionemos todas las columnas menos la columna team:

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

Condiciones complejas con .selectExpr()

Si manejamos bien SQL y tenemos que aplicar condiciones más complejas al filtrar columnas, el método .selectExpr() de PySpark ofrece una potente solución. Nos permite utilizar expresiones similares a SQL para seleccionar y manipular columnas directamente en nuestro código PySpark.

Veamos este ejemplo:

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

Aquí no solo excluimos una columna. Estamos haciendo lo siguiente:

1. Seleccionar las columnas player_name y player_position.
2. Crear una nueva columna llamada FullMatch que indica si un jugador ha jugado un partido completo (60 minutos o más).

Esta flexibilidad hace que .selectExpr() tenga un valor incalculable para situaciones en las que las simples comprensiones de listas no sean suficientes.

Prácticas recomendadas

Cuando decidamos eliminar columnas de nuestro DataFrame, seguir algunas prácticas recomendadas es crucial para mantener la integridad de los datos, la legibilidad del código y la eficiencia.

1. Haz una copia de seguridad de tus datos: Antes de realizar cualquier modificación en tu DataFrame, especialmente al eliminar columnas, es conveniente crear una copia de seguridad. Esto garantiza que puedas volver a los datos originales si es necesario.

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

2. Elimina con inplace=False (de forma predeterminada): De forma predeterminada, el método .drop() devuelve un nuevo DataFrame sin modificar el original. En general, esto es más seguro. Si no tienes ninguna duda de que quieres modificar el DataFrame original, puedes utilizar inplace=True, pero ten cuidado.

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

3. Documenta tus decisiones: Si estás trabajando en un proyecto colaborativo o prevés que volverás a usar tu código en el futuro, es útil añadir comentarios que expliquen por qué has decidido eliminar determinadas columnas. Esta documentación puede proporcionar un contexto valioso para ti y para los demás.

4. Considera la posibilidad de cambiar el nombre de las columnas: En lugar de eliminar columnas, debemos plantearnos si cambiarles el nombre sería más adecuado, sobre todo si la información que contienen puede ser valiosa para futuros análisis. PySpark nos proporciona el método .withColumnRenamed(), que nos ayuda a cambiar el nombre de las columnas.

Conclusión

En este tutorial, hemos aprendido a eliminar una o más columnas utilizando los métodos drop() y .select(). También hemos descrito métodos alternativos para aprovechar las expresiones SQL si necesitamos técnicas de selección más complejas.

Si quieres ampliar tus conocimientos y aprovechar las potentes técnicas de PySpark, considera la posibilidad de echar un vistazo a estos cursos y materiales: 

• Big data con PySpark
• Fundamentos de big data con PySpark
• Machine learning con PySpark