Seleccione filas de un DataFrame basado en valores en una columna en pandas

¿Cómo seleccionar filas de un DataFrame basado en valores en alguna columna en pandas?
En SQL yo usaría:

select * from table where colume_name = some_value. 

Intenté ver la documentación de los pandas, pero no encontré la respuesta de inmediato.

Para seleccionar filas cuyo valor de columna sea igual a un escalar, some_value , use == :

 df.loc[df['column_name'] == some_value] 

Para seleccionar filas cuyo valor de columna esté en un iterable, some_values , use isin :

 df.loc[df['column_name'].isin(some_values)] 

Combina múltiples condiciones con & :

 df.loc[(df['column_name'] >= A) & (df['column_name'] <= B)] 

Tenga en cuenta los paréntesis. Debido a las reglas de precedencia del operador de Python, & enlaza más estrechamente que <= y >= . Por lo tanto, los paréntesis en el último ejemplo son necesarios. Sin los paréntesis

 df['column_name'] >= A & df['column_name'] <= B 

se analiza como

 df['column_name'] >= (A & df['column_name']) <= B 

lo que resulta en un valor de Verdad de una Serie es un error ambiguo .


Para seleccionar filas cuyo valor de columna no sea igual a some_value , use != :

 df.loc[df['column_name'] != some_value] 

isin devuelve una serie booleana, por lo que para seleccionar filas cuyo valor no está en some_values , some_values la serie booleana utilizando ~ :

 df.loc[~df['column_name'].isin(some_values)] 

Por ejemplo,

 import pandas as pd import numpy as np df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(), 'B': 'one one two three two two one three'.split(), 'C': np.arange(8), 'D': np.arange(8) * 2}) print(df) # ABCD # 0 foo one 0 0 # 1 bar one 1 2 # 2 foo two 2 4 # 3 bar three 3 6 # 4 foo two 4 8 # 5 bar two 5 10 # 6 foo one 6 12 # 7 foo three 7 14 print(df.loc[df['A'] == 'foo']) 

rendimientos

  ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Si tiene varios valores que desea incluir, póngalos en una lista (o más generalmente, cualquier iterable) y use isin :

 print(df.loc[df['B'].isin(['one','three'])]) 

rendimientos

  ABCD 0 foo one 0 0 1 bar one 1 2 3 bar three 3 6 6 foo one 6 12 7 foo three 7 14 

Sin embargo, tenga en cuenta que si desea hacer esto muchas veces, es más eficiente hacer un índice primero y luego usar df.loc :

 df = df.set_index(['B']) print(df.loc['one']) 

rendimientos

  ACD B one foo 0 0 one bar 1 2 one foo 6 12 

o, para incluir múltiples valores del índice use df.index.isin :

 df.loc[df.index.isin(['one','two'])] 

rendimientos

  ACD B one foo 0 0 one bar 1 2 two foo 2 4 two foo 4 8 two bar 5 10 one foo 6 12 

tl; dr

Los pandas equivalen a

 select * from table where column_name = some_value 

es

 table[table.column_name == some_value] 

Múltiples condiciones:

 table[(table.column_name == some_value) | (table.column_name2 == some_value2)] 

o

 table.query('column_name == some_value | column_name2 == some_value2') 

Ejemplo de código

 import pandas as pd # Create data set d = {'foo':[100, 111, 222], 'bar':[333, 444, 555]} df = pd.DataFrame(d) # Full dataframe: df # Shows: # bar foo # 0 333 100 # 1 444 111 # 2 555 222 # Output only the row(s) in df where foo is 222: df[df.foo == 222] # Shows: # bar foo # 2 555 222 

En el código anterior, es la línea df[df.foo == 222] que da las filas basadas en el valor de la columna, 222 en este caso.

Múltiples condiciones también son posibles:

 df[(df.foo == 222) | (df.bar == 444)] # bar foo # 1 444 111 # 2 555 222 

Pero en ese momento recomendaría usar la función de consulta , ya que es menos detallada y produce el mismo resultado:

 df.query('foo == 222 | bar == 444') 

Hay algunas formas básicas de seleccionar filas de un dataframe de pandas.

  1. Indexación booleana
  2. Indexación posicional
  3. Indexación de tags
  4. API

Para cada tipo de base, podemos mantener las cosas simples al restringirnos a la API de pandas o podemos aventurarnos fuera de la API, por lo general en numpy , y acelerar las cosas.

Le mostraré ejemplos de cada uno y le guiaré para saber cuándo usar ciertas técnicas.


Preparar
Lo primero que necesitaremos es identificar una condición que actuará como nuestro criterio para seleccionar filas. El OP ofrece hasta column_name == some_value . Comenzaremos allí e incluiremos algunos otros casos de uso comunes.

Préstamo de @unutbu:

 import pandas as pd, numpy as np df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(), 'B': 'one one two three two two one three'.split(), 'C': np.arange(8), 'D': np.arange(8) * 2}) 

Supongamos que nuestro criterio es la columna 'A' = 'foo'

1.
La indexación booleana requiere encontrar el valor verdadero de la columna 'A' de cada fila que sea igual a 'foo' , y luego usar esos valores de verdad para identificar qué filas mantener. Normalmente, nombraríamos esta serie, una matriz de valores de verdad, mask . Lo haremos aquí también.

 mask = df['A'] == 'foo' 

Luego podemos usar esta máscara para cortar o indexar el dataframe.

 df[mask] ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Esta es una de las formas más sencillas de realizar esta tarea y si el rendimiento o la intuición no son un problema, este debería ser su método elegido. Sin embargo, si el rendimiento es una preocupación, es posible que desee considerar una forma alternativa de crear la mask .


2.
La indexación posicional tiene sus casos de uso, pero este no es uno de ellos. Para identificar dónde dividir, primero debemos realizar el mismo análisis booleano que hicimos anteriormente. Esto nos deja realizando un paso adicional para realizar la misma tarea.

 mask = df['A'] == 'foo' pos = np.flatnonzero(mask) df.iloc[pos] ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

3.
La indexación de tags puede ser muy útil, pero en este caso, nuevamente estamos haciendo más trabajo sin beneficio

 df.set_index('A', append=True, drop=False).xs('foo', level=1) ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

4.
pd.DataFrame.query es una forma muy elegante / intuitiva de realizar esta tarea. Pero a menudo es más lento. Sin embargo , si presta atención a los horarios a continuación, para datos grandes, la consulta es muy eficiente. Más que el enfoque estándar y de magnitud similar a mi mejor sugerencia.

 df.query('A == "foo"') ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Mi preferencia es usar la mask Boolean

Se pueden realizar mejoras reales modificando cómo creamos nuestra mask Boolean .

mask alternativa 1
Utilice la matriz numpy subyacente y renuncie a la sobrecarga de crear otro pd.Series

 mask = df['A'].values == 'foo' 

Mostraré más pruebas de tiempo completas al final, pero solo eche un vistazo a las mejoras de rendimiento que obtenemos utilizando el dataframe de muestra. Primero, observamos la diferencia en la creación de la mask

 %timeit mask = df['A'].values == 'foo' %timeit mask = df['A'] == 'foo' 5.84 µs ± 195 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) 166 µs ± 4.45 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) 

La evaluación de la mask con la matriz numpy es ~ 30 veces más rápida. Esto se debe en parte a que la evaluación de numpy menudo más rápida. También se debe en parte a la falta de gastos generales necesarios para crear un índice y un objeto pd.Series correspondiente.

A continuación, veremos el momento para cortar con una mask comparación con la otra.

 mask = df['A'].values == 'foo' %timeit df[mask] mask = df['A'] == 'foo' %timeit df[mask] 219 µs ± 12.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) 239 µs ± 7.03 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) 

Las ganancias de rendimiento no son tan pronunciadas. Veremos si esto se mantiene sobre pruebas más robustas.


mask alternativa 2
Podríamos haber reconstruido el dataframe también. Hay una gran advertencia al reconstruir un dataframe: ¡debe tener cuidado con los dtypes al hacerlo!

En lugar de df[mask] haremos esto

 pd.DataFrame(df.values[mask], df.index[mask], df.columns).astype(df.dtypes) 

Si el dataframe es de tipo mixto, que es nuestro ejemplo, entonces cuando obtengamos df.values la matriz resultante es de dtype object y, en consecuencia, todas las columnas del nuevo dataframe serán de dtype object . Por lo tanto, se requiere el astype(df.dtypes) y se astype(df.dtypes) las posibles ganancias de rendimiento.

 %timeit df[m] %timeit pd.DataFrame(df.values[mask], df.index[mask], df.columns).astype(df.dtypes) 216 µs ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) 1.43 ms ± 39.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) 

Sin embargo, si el dataframe no es de tipo mixto, esta es una forma muy útil de hacerlo.

Dado

 np.random.seed([3,1415]) d1 = pd.DataFrame(np.random.randint(10, size=(10, 5)), columns=list('ABCDE')) d1 ABCDE 0 0 2 7 3 8 1 7 0 6 8 6 2 0 2 0 4 9 3 7 3 2 4 3 4 3 6 7 7 4 5 5 3 7 5 9 6 8 7 6 4 7 7 6 2 6 6 5 8 2 8 7 5 8 9 4 7 6 1 5 

 %%timeit mask = d1['A'].values == 7 d1[mask] 179 µs ± 8.73 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) 

Versus

 %%timeit mask = d1['A'].values == 7 pd.DataFrame(d1.values[mask], d1.index[mask], d1.columns) 87 µs ± 5.12 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) 

Cortamos el tiempo a la mitad.


mask alternativa 3
@unutbu también nos muestra cómo usar pd.Series.isin para dar cuenta de que cada elemento de df['A'] se encuentra en un conjunto de valores. Esto se evalúa de la misma manera si nuestro conjunto de valores es un conjunto de un valor, a saber, 'foo' . Pero también se generaliza para incluir conjuntos más grandes de valores si es necesario. Resulta que esto sigue siendo bastante rápido, aunque es una solución más general. La única pérdida real es la intuición para quienes no están familiarizados con el concepto.

 mask = df['A'].isin(['foo']) df[mask] ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Sin embargo, al igual que antes, podemos utilizar el numpy para mejorar el rendimiento sin sacrificar prácticamente nada. Usaremos np.in1d

 mask = np.in1d(df['A'].values, ['foo']) df[mask] ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Sincronización
Incluiré otros conceptos mencionados en otras publicaciones como referencia.
Código abajo

Cada columna en esta tabla representa un dataframe de longitud diferente sobre el cual probamos cada función. Cada columna muestra el tiempo relativo tomado, con la función más rápida dado un índice base de 1.0 .

 res.div(res.min()) 10 30 100 300 1000 3000 10000 30000 mask_standard 2.156872 1.850663 2.034149 2.166312 2.164541 3.090372 2.981326 3.131151 mask_standard_loc 1.879035 1.782366 1.988823 2.338112 2.361391 3.036131 2.998112 2.990103 mask_with_values 1.010166 1.000000 1.005113 1.026363 1.028698 1.293741 1.007824 1.016919 mask_with_values_loc 1.196843 1.300228 1.000000 1.000000 1.038989 1.219233 1.037020 1.000000 query 4.997304 4.765554 5.934096 4.500559 2.997924 2.397013 1.680447 1.398190 xs_label 4.124597 4.272363 5.596152 4.295331 4.676591 5.710680 6.032809 8.950255 mask_with_isin 1.674055 1.679935 1.847972 1.724183 1.345111 1.405231 1.253554 1.264760 mask_with_in1d 1.000000 1.083807 1.220493 1.101929 1.000000 1.000000 1.000000 1.144175 

Notará que los tiempos más rápidos parecen compartirse entre mask_with_values y mask_with_in1d

 res.T.plot(loglog=True) 

introduzca la descripción de la imagen aquí

Funciones

 def mask_standard(df): mask = df['A'] == 'foo' return df[mask] def mask_standard_loc(df): mask = df['A'] == 'foo' return df.loc[mask] def mask_with_values(df): mask = df['A'].values == 'foo' return df[mask] def mask_with_values_loc(df): mask = df['A'].values == 'foo' return df.loc[mask] def query(df): return df.query('A == "foo"') def xs_label(df): return df.set_index('A', append=True, drop=False).xs('foo', level=-1) def mask_with_isin(df): mask = df['A'].isin(['foo']) return df[mask] def mask_with_in1d(df): mask = np.in1d(df['A'].values, ['foo']) return df[mask] 

Pruebas

 res = pd.DataFrame( index=[ 'mask_standard', 'mask_standard_loc', 'mask_with_values', 'mask_with_values_loc', 'query', 'xs_label', 'mask_with_isin', 'mask_with_in1d' ], columns=[10, 30, 100, 300, 1000, 3000, 10000, 30000], dtype=float ) for j in res.columns: d = pd.concat([df] * j, ignore_index=True) for i in res.index:a stmt = '{}(d)'.format(i) setp = 'from __main__ import d, {}'.format(i) res.at[i, j] = timeit(stmt, setp, number=50) 

Tiempo especial
Observando el caso especial cuando tenemos un único tipo de objeto dtype para todo el dataframe. Código abajo

 spec.div(spec.min()) 10 30 100 300 1000 3000 10000 30000 mask_with_values 1.009030 1.000000 1.194276 1.000000 1.236892 1.095343 1.000000 1.000000 mask_with_in1d 1.104638 1.094524 1.156930 1.072094 1.000000 1.000000 1.040043 1.027100 reconstruct 1.000000 1.142838 1.000000 1.355440 1.650270 2.222181 2.294913 3.406735 

Resulta que la reconstrucción no vale la pena después de unos pocos cientos de filas.

 spec.T.plot(loglog=True) 

introduzca la descripción de la imagen aquí

Funciones

 np.random.seed([3,1415]) d1 = pd.DataFrame(np.random.randint(10, size=(10, 5)), columns=list('ABCDE')) def mask_with_values(df): mask = df['A'].values == 'foo' return df[mask] def mask_with_in1d(df): mask = np.in1d(df['A'].values, ['foo']) return df[mask] def reconstruct(df): v = df.values mask = np.in1d(df['A'].values, ['foo']) return pd.DataFrame(v[mask], df.index[mask], df.columns) spec = pd.DataFrame( index=['mask_with_values', 'mask_with_in1d', 'reconstruct'], columns=[10, 30, 100, 300, 1000, 3000, 10000, 30000], dtype=float ) 

Pruebas

 for j in spec.columns: d = pd.concat([df] * j, ignore_index=True) for i in spec.index: stmt = '{}(d)'.format(i) setp = 'from __main__ import d, {}'.format(i) spec.at[i, j] = timeit(stmt, setp, number=50) 

Encuentro que la syntax de las respuestas anteriores es redundante y difícil de recordar. Pandas introdujo el método query() en v0.13 y lo prefiero mucho. Para su pregunta, podría hacer df.query('col == val')

Reproducido de http://pandas.pydata.org/pandas-docs/version/0.17.0/indexing.html#indexing-query

 In [167]: n = 10 In [168]: df = pd.DataFrame(np.random.rand(n, 3), columns=list('abc')) In [169]: df Out[169]: abc 0 0.687704 0.582314 0.281645 1 0.250846 0.610021 0.420121 2 0.624328 0.401816 0.932146 3 0.011763 0.022921 0.244186 4 0.590198 0.325680 0.890392 5 0.598892 0.296424 0.007312 6 0.634625 0.803069 0.123872 7 0.924168 0.325076 0.303746 8 0.116822 0.364564 0.454607 9 0.986142 0.751953 0.561512 # pure python In [170]: df[(df.a < df.b) & (df.b < df.c)] Out[170]: abc 3 0.011763 0.022921 0.244186 8 0.116822 0.364564 0.454607 # query In [171]: df.query('(a < b) & (b < c)') Out[171]: abc 3 0.011763 0.022921 0.244186 8 0.116822 0.364564 0.454607 

También puede acceder a las variables en el entorno al anteponer una @ .

 exclude = ('red', 'orange') df.query('color not in @exclude') 

Se pueden lograr resultados más rápidos usando numpy.where .

Por ejemplo, con la configuración de Unubtu –

 In [76]: df.iloc[np.where(df.A.values=='foo')] Out[76]: ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Comparaciones de tiempo:

 In [68]: %timeit df.iloc[np.where(df.A.values=='foo')] # fastest 1000 loops, best of 3: 380 µs per loop In [69]: %timeit df.loc[df['A'] == 'foo'] 1000 loops, best of 3: 745 µs per loop In [71]: %timeit df.loc[df['A'].isin(['foo'])] 1000 loops, best of 3: 562 µs per loop In [72]: %timeit df[df.A=='foo'] 1000 loops, best of 3: 796 µs per loop In [74]: %timeit df.query('(A=="foo")') # slowest 1000 loops, best of 3: 1.71 ms per loop 

Aquí hay un ejemplo simple.

 from pandas import DataFrame # Create data set d = {'Revenue':[100,111,222], 'Cost':[333,444,555]} df = DataFrame(d) # mask = Return True when the value in column "Revenue" is equal to 111 mask = df['Revenue'] == 111 print mask # Result: # 0 False # 1 True # 2 False # Name: Revenue, dtype: bool # Select * FROM df WHERE Revenue = 111 df[mask] # Result: # Cost Revenue # 1 444 111 

Acabo de intentar editar esto, pero no he iniciado sesión, así que no estoy seguro de a dónde fue mi edición. Estaba tratando de incorporar selección múltiple. Así que creo que una mejor respuesta es:

Para un solo valor, el más directo (legible por humanos) es probablemente:

 df.loc[df['column_name'] == some_value] 

Para las listas de valores también puede utilizar:

 df.loc[df['column_name'].isin(some_values)] 

Por ejemplo,

 import pandas as pd import numpy as np df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(), 'B': 'one one two three two two one three'.split(), 'C': np.arange(8), 'D': np.arange(8) * 2}) print(df) # ABCD # 0 foo one 0 0 # 1 bar one 1 2 # 2 foo two 2 4 # 3 bar three 3 6 # 4 foo two 4 8 # 5 bar two 5 10 # 6 foo one 6 12 # 7 foo three 7 14 print(df.loc[df['A'] == 'foo']) 

rendimientos

  ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Si tiene varios criterios contra los que desea seleccionar, puede ponerlos en una lista y usar ‘isin’:

 print(df.loc[df['B'].isin(['one','three'])]) 

rendimientos

  ABCD 0 foo one 0 0 1 bar one 1 2 3 bar three 3 6 6 foo one 6 12 7 foo three 7 14 

Sin embargo, tenga en cuenta que si desea hacer esto muchas veces, es más eficiente hacer que A sea el índice primero y luego usar df.loc:

 df = df.set_index(['A']) print(df.loc['foo']) 

rendimientos

  ABCD foo one 0 0 foo two 2 4 foo two 4 8 foo one 6 12 foo three 7 14 

Si encuentra filas basadas en algún entero en una columna, entonces

 df.loc[df['column_name'] == 2017] 

Si estás encontrando valor basado en cadena

 df.loc[df['column_name'] == 'string'] 

Si se basa en ambos

 df.loc[(df['column_name'] == 'string') & (df['column_name'] == 2017)] 
 df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(), 'B': 'one one two three two two one three'.split(), 'C': np.arange(8), 'D': np.arange(8) * 2}) df[df['A']=='foo'] OUTPUT: ABCD 0 foo one 0 0 2 foo two 2 4 4 foo two 4 8 6 foo one 6 12 7 foo three 7 14 

Para adjuntar a esta famosa pregunta (aunque un poco demasiado tarde): También puede hacer df.groupby('column_name').get_group('column_desired_value').reset_index() para crear un nuevo dataframe con una columna específica que tenga un valor particular . P.ej

 import pandas as pd df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(), 'B': 'one one two three two two one three'.split()}) print("Original dataframe:") print(df) b_is_two_dataframe = pd.DataFrame(df.groupby('B').get_group('two').reset_index()).drop('index', axis = 1) #NOTE: the final drop is to remove the extra index column returned by groupby object print('Sub dataframe where B is two:') print(b_is_two_dataframe) 

Ejecutar esto da:

 Original dataframe: AB 0 foo one 1 bar one 2 foo two 3 bar three 4 foo two 5 bar two 6 foo one 7 foo three Sub dataframe where B is two: AB 0 foo two 1 foo two 2 bar two 

Para seleccionar solo columnas específicas de varias columnas para un valor dado en pandas:

 select col_name1, col_name2 from table where column_name = some_value. 

Opciones:

 df.loc[df['column_name'] == some_value][[col_name1, col_name2]] 

o

 df.query['column_name' == 'some_value'][[col_name1, col_name2]] 

Si vino aquí para seleccionar filas de un dataframe incluyendo aquellos cuyo valor de columna NO es cualquiera de una lista de valores, a continuación se muestra cómo voltear la respuesta de unutbu para obtener una lista de valores arriba:

 df.loc[~df['column_name'].isin(some_values)] 

(Para no incluir un solo valor, por supuesto, solo usa el operador regular no es igual a,! != .)

Ejemplo:

 import pandas as pd df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(), 'B': 'one one two three two two one three'.split()}) print(df) 

Nos da

  AB 0 foo one 1 bar one 2 foo two 3 bar three 4 foo two 5 bar two 6 foo one 7 foo three 

Para agrupar solo esas filas que NO ESTÁN one o three en la columna B :

 df.loc[~df['B'].isin(['one', 'three'])] 

rendimientos

  AB 2 foo two 4 foo two 5 bar two 

También puedes usar .apply:

 df.apply(lambda row: row[df['B'].isin(['one','three'])]) 

En realidad, funciona en forma de fila (es decir, aplica la función a cada fila).

La salida es

  ABCD 0 foo one 0 0 1 bar one 1 2 3 bar three 3 6 6 foo one 6 12 7 foo three 7 14 

Los resultados son los mismos que se utilizan como se menciona en @unutbu

 df[[df['B'].isin(['one','three'])]] 
 df.loc[df['column_name'] == some_value]