Uso de Pandas Query para filtrar datos de forma sencilla

Publicado el 02 mayo 2022 por Daniel Rodríguez @analyticslane

En los DataFrames de Pandas se puede emplear vectores lógicos para seleccionar las filas que cumplen una o varias condiciones. Algo que se puede utilizar para realizar operaciones de filtrado de datos. Aunque esta característica es muy poderosa, existe otra opción con la que generalmente se puede conseguir filtrar los datos de una manera mucho más sencilla: Pandas Query. Un método disponible en los DataFrames de Pandas que evalúa una condición lógica y devuelve el conjunto de datos que la cumple.

Selección tradicional de datos en Pandas

Tal como se ha explicado en más detalle anteriormente, es posible seleccionar las filas de un DataFrame utilizando vectores con valores lógicos. Algo que también se puede hacer cuando los valores son cadenas de texto. Con lo que se puede filtrar valores del conjunto de datos. Así, si se realiza una comparación sobre una columna se pueden seleccionar los datos que la cumplen. Por ejemplo, para seleccionar los exoplanetas del conjunto de datos planets que tienen una masa dada solamente se tiene que ejecutar el siguiente código.

import seaborn as sb

planets = sb.load_dataset('planets')
planets.head()

planets[planets.mass == 10.5]
            method  number  orbital_period  mass  distance  year
4  Radial Velocity       1          516.22  10.5    119.47  2009

Obteniendo como resultado el único planeta con una masa 10.5 la de la tierra

Selección de datos con Pandas Query

Esta misma operación de filtrado se puede realizar de una manera más sencilla con Pandas Query. Para ello solamente se tienen que llamar al método e incluir como parámetro una cadena de texto con la condición, usando el nombre de la columna para identificar aquella sobre la que se realiza la operación. Esto es, escribiendo la siguiente línea de código.

planets.query('mass == 10.5')
            method  number  orbital_period  mass  distance  year
4  Radial Velocity       1          516.22  10.5    119.47  2009

Con lo que se ha obtenido exactamente el mismo resultado que antes. Aunque Pandas Query ofrece muchas más posibilidades.

Uso de varias condiciones

Es posible combinar más de una condición mediante el uso de los diferentes operadores lógicos. Así se pueden seleccionar los exoplanetas descubiertos en 2010 que tengan una masa inferior a 5% de la masa terrestre.

planets.query('year == 2010 and mass < 0.05')
              method  number  orbital_period   mass  distance  year
145  Radial Velocity       4       124.26000  0.046      4.70  2010
149  Radial Velocity       1         2.64561  0.022     19.80  2010
230  Radial Velocity       6         5.75962  0.041     39.39  2010
231  Radial Velocity       6        16.35670  0.038     39.39  2010
399  Radial Velocity       1         9.49400  0.026     21.30  2010
506  Radial Velocity       1         4.64550  0.013     24.05  2010

Como se puede ver simplemente se han escrito las dos condiciones unidad por el operador lógico and.

Comparar cadenas de texto

Los valores de las columnas que se evalúan no tienen por qué ser únicamente numéricas, también se puede comparar con cadenas de texto. Para lo que únicamente se tienen que indicar el valor de la cadena a comparar entrecomillado. De este modo, para obtener los exoplanetas que han descubierto mediante el método Astrometry solamente se tienen que escribir algo como la siguiente línea.

planets.query('method == "Astrometry"')
         method  number  orbital_period  mass  distance  year
113  Astrometry       1          246.36   NaN     20.77  2013
537  Astrometry       1         1016.00   NaN     14.98  2010

Obteniendo como resultado los dos exoplanetas que aparecen en el conjunto de datos.

Uso de variables en las comparaciones

Cuando el valor que se desea comparar se encuentra en una variable no es necesario construir la cadena de texto. El método query() puede evaluar el valor simplemente escribiendo el nombre de la variable precedida de una @. Algo que permite usar el método no sólo para análisis estáticos. Así, para seleccionar los exoplanetas que cuyo método de descubrimiento se ha indicado en una variable se puede emplear algo como lo que se muestra a continuación.

method = 'Orbital Brightness Modulation'

planets.query('method == @method')
                            method  number  orbital_period  mass  distance  year
787  Orbital Brightness Modulation       2        0.240104   NaN    1180.0  2011
788  Orbital Brightness Modulation       2        0.342887   NaN    1180.0  2011
792  Orbital Brightness Modulation       1        1.544929   NaN       NaN  2013

Valores en una lista

Otra posibilidad es seleccionar los valores en una lista gracias al operador in. Por ejemplo, se puede obtener los planetas que se han descubierto con una lista de métodos.

planets.query('method in ["Astrometry", "Orbital Brightness Modulation"]')
                            method  number  orbital_period  mass  distance  year
113                     Astrometry       1      246.360000   NaN     20.77  2013
537                     Astrometry       1     1016.000000   NaN     14.98  2010
787  Orbital Brightness Modulation       2        0.240104   NaN   1180.00  2011
788  Orbital Brightness Modulation       2        0.342887   NaN   1180.00  2011
792  Orbital Brightness Modulation       1        1.544929   NaN       NaN  2013

Lista que obviamente puede estar en una variable. El siguiente código generaría exactamente el mismo resultado.

methods = ["Astrometry", "Orbital Brightness Modulation"]

planets.query('method in @methods')

Identificar valores nulos y no nulos

Para poder identificar los valores que son nulos (o no nulos) es necesario recurrir a un truco. Simplemente se ha de comparar una columna con sigo misma. En los registros que el valor sea nulo esta comparación es falsa, mientras que en resto será verdadera. Esto es algo que se puede ver mejor con un ejemplo. Como se ha visto en la última consulta, existen dos exoplanetas descubiertos por el método Orbital Brightness Modulation de los que se conoce la distancia y uno que no, para seleccionar aquel en el que no se conoce la distancia simplemente se debe escribir.

planets.query('distance != distance and method == "Orbital Brightness Modulation"')
                            method  number  orbital_period  mass  distance  year
792  Orbital Brightness Modulation       1        1.544929   NaN       NaN  2013

Por otro lado, para seleccionar aquellos en los que la distancia es conocida se debería escribir.

planets.query('distance == distance and method == "Orbital Brightness Modulation"')
                            method  number  orbital_period  mass  distance  year
787  Orbital Brightness Modulation       2        0.240104   NaN    1180.0  2011
788  Orbital Brightness Modulation       2        0.342887   NaN    1180.0  2011

Valores con cadenas de texto

Para seleccionar los valores que contengan dentro de la cadena un texto es necesario recurrir a los métodos de las clases de Python. Siendo necesario indicarle al método que use Python como motor de ejecución. Así se puede usar str.contains() para localizar los métodos que contienen la cadena Variations.

planets.query('method.str.contains("Variations")', engine='python')
                          method  number  orbital_period  mass  distance  year
32     Eclipse Timing Variations       1      10220.0000  6.05       NaN  2009
37     Eclipse Timing Variations       2       5767.0000   NaN    130.72  2008
38     Eclipse Timing Variations       2       3321.0000   NaN    130.72  2008
39     Eclipse Timing Variations       2       5573.5500   NaN    500.00  2010
40     Eclipse Timing Variations       2       2883.5000   NaN    500.00  2010
41     Eclipse Timing Variations       1       2900.0000   NaN       NaN  2011
42     Eclipse Timing Variations       1       4343.5000  4.20       NaN  2012
43     Eclipse Timing Variations       2       5840.0000   NaN       NaN  2011
44     Eclipse Timing Variations       2       1916.2500   NaN       NaN  2011
680    Transit Timing Variations       2        160.0000   NaN   2119.00  2011
736    Transit Timing Variations       2         57.0110   NaN    855.00  2012
749    Transit Timing Variations       3             NaN   NaN       NaN  2014
813    Transit Timing Variations       2         22.3395   NaN    339.00  2013
958  Pulsation Timing Variations       1       1170.0000   NaN       NaN  2007

Nombres de columnas con espacios

Aunque no es el caso del ejemplo, los nombres de las columnas de los DataFrame de Pandas admiten nombres con espacios en blanco. Algo que no es recomendable. Para filtrar en base a estas columnas tiene que envolver el nombre entre comillas invertidas ("). Así, si se cambia el nombre de la columna orbital_period por orbital period se puede seguir usando Pandas Query sin mayor problema.

planets.columns = ['method', 'number', 'orbital period', 'mass', 'distance', 'year']

planets.query('`orbital period` < 0.5')
                            method  number  orbital period  mass  distance  year
730                        Transit       3        0.453285   NaN      38.7  2011
787  Orbital Brightness Modulation       2        0.240104   NaN    1180.0  2011
788  Orbital Brightness Modulation       2        0.342887   NaN    1180.0  2011
794                        Transit       1        0.355000   NaN       NaN  2013
945                  Pulsar Timing       1        0.090706   NaN    1200.0  2011

Conclusiones

Hoy hemos visto un método disponible en los DataFrame Pandas para filtrar los datos de forma sencilla. Con lo que estas operaciones tan habituales pueden realizarse de una forma eficiente. ¿Opinas que Pandas Query es más sencillo o no? Puedes dejar tus opiniones en los comentarios.