Filtrado y Selección de Datos con Python

En la mayoría de los escenarios reales, no es necesario trabajar con todas las columnas o filas de un conjunto de datos. Las tablas grandes pueden tener cientos de columnas y millones de filas, la mayoría de las cuales pueden ser irrelevantes para un análisis específico. En este capítulo, aprenderás cómo "extraer" exactamente los datos que necesitas utilizando los métodos select() y filter().

Selección de columnas específicas

El método select() permite crear un nuevo DataFrame que solo contiene las columnas que elijas. Esto reduce la cantidad de memoria que utiliza tu clúster y facilita la lectura de los resultados.

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

display(selected_df)

Observa que Spark no modifica el df original. En su lugar, crea uno nuevo llamado selected_df. Si se quisiera seleccionar todas las columnas excepto una, o realizar un cálculo durante la selección, se utilizaría una sintaxis más avanzada, pero para tareas básicas, pasar los nombres de las columnas como cadenas es el enfoque estándar.

Filtrado de filas con condiciones

El método filter() (o su alias where()) actúa como un filtro para los datos. Se proporciona una condición y Spark solo conserva las filas donde esa condición es verdadera.

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

display(filtered_df)

Se pueden utilizar operadores de comparación estándar como == (igual), != (diferente), > (mayor que) y < (menor que).

Combinación de múltiples filtros

A menudo es necesario aplicar más de una regla al mismo tiempo. Para ello, se pueden encadenar filtros o utilizar operadores lógicos como "And" (&) y "Or" (|).

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

display(good_carat_df)

Selección y filtrado en un solo paso

Debido a que Spark utiliza una API "fluida", es posible encadenar estos comandos en una sola línea de código. Este es un patrón muy común en la ingeniería de datos profesional:

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

display(high_profit_df)

Comprobación del trabajo

Después de cada selección o filtro, es una buena práctica ejecutar un count(). Si comienzas con 10,000 filas y después de un filtro tienes 0, puedes saber que tu condición de filtro podría ser demasiado estricta o contener un error tipográfico en los valores de cadena.