Pregunta Seleccionar 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 usaría:

select * from table where colume_name = some_value. 

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


821
2018-06-12 17:42


origen


Respuestas:


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

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

Para seleccionar filas cuyo valor de columna está en un iterable, some_values, utilizar isin:

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

Combina múltiples condiciones con &:

df.loc[(df['column_name'] == some_value) & df['other_column'].isin(some_values)]

Para seleccionar filas cuyo valor de columna no es igual  some_value, utilizar !=:

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

isin devuelve una serie booleana, para seleccionar filas cuyo valor es no en some_values, niegue la Serie booleana usando ~:

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)
#      A      B  C   D
# 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

     A      B  C   D
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, colóquelos en un lista (o más generalmente, cualquier iterable) y uso isin:

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

rendimientos

     A      B  C   D
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

Tenga en cuenta, sin embargo, 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

       A  C   D
B              
one  foo  0   0
one  bar  1   2
one  foo  6  12

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

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

rendimientos

       A  C   D
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

1646
2018-06-12 17:44



tl; dr

Los pandas equivalentes 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 punto recomendaría usar consulta función, ya que es menos detallado y produce el mismo resultado:

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

134
2017-07-08 15:17



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

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

Para cada tipo de base, podemos mantener las cosas simples restringiéndonos a la API de pandas o podemos aventurarnos fuera de la API, generalmente en numpyy acelerar las cosas

Te mostraré ejemplos de cada uno y te guiaré sobre 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 común.

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 columna 'A' = 'foo'

1.
Booleano la indexación requiere encontrar el valor de verdad de cada fila 'A' la columna es igual a 'foo', luego usa esos valores de verdad para identificar qué filas mantener. Normalmente, le daríamos el nombre a esta serie, una matriz de valores de verdad, mask. Lo haremos aquí también.

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

Podemos usar esta máscara para cortar o indexar el marco de datos

df[mask]

     A      B  C   D
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 simples de llevar a cabo esta tarea y si el rendimiento o la intuición no es un problema, este debería ser su método elegido. Sin embargo, si el rendimiento es una preocupación, entonces es posible que desee considerar una forma alternativa de crear el mask.


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

mask = df['A'] == 'foo'
pos = np.flatnonzero(mask)
df.iloc[pos]

     A      B  C   D
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.
Etiqueta indexar puede ser muy útil, pero en este caso, nuevamente estamos haciendo más trabajo sin ningún beneficio

df.set_index('A', append=True, drop=False).xs('foo', level=1)

     A      B  C   D
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 tiempos que se muestran a continuación, para los datos grandes, la consulta es muy eficiente. Más que el enfoque estándar y de una magnitud similar a mi mejor sugerencia.

df.query('A == "foo"')

     A      B  C   D
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 el Boolean  mask 

Se pueden hacer mejoras reales modificando la forma en que creamos nuestro Boolean  mask.

mask alternativa 1
Usa el subyacente numpy matriz y renunciar a la sobrecarga de la creación de otro pd.Series 

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

Mostraré pruebas de tiempo más completas al final, pero solo eche un vistazo a las ganancias de rendimiento que obtenemos con el marco de datos de muestra. Primero vemos la diferencia en la creación del 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)

Evaluar el mask con el numpy la matriz es ~ 30 veces más rápida. Esto se debe en parte a numpy la evaluación a menudo es más rápida. También se debe en parte a la falta de sobrecarga necesaria para construir un índice y un correspondiente pd.Series objeto.

A continuación veremos el momento para cortar con una mask contra el otro.

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 sostiene sobre pruebas más robustas.


mask alternativa 2
Podríamos haber reconstruido el marco de datos también. Al reconstruir un marco de datos hay una gran advertencia: debe encargarse de dtypes cuando lo haces!

En lugar de df[mask] nosotros haremos esto

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

Si el marco de datos es de tipo mixto, que es nuestro ejemplo, cuando lo logremos df.values la matriz resultante es de dtype  object y, en consecuencia, todas las columnas del nuevo marco de datos serán de dtype  object. Por lo tanto, requiere el astype(df.dtypes) y matando cualquier ganancia de rendimiento potencial.

%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

   A  B  C  D  E
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 cada elemento de df['A'] estar en un conjunto de valores. Esto se evalúa de la misma manera si nuestro conjunto de valores es un conjunto de un valor, es decir 'foo'. Pero también se generaliza para incluir conjuntos de valores más grandes si es necesario. Resulta que esto todavía es bastante rápido a pesar de que es una solución más general. La única pérdida real es en la intuición para aquellos que no están familiarizados con el concepto.

mask = df['A'].isin(['foo'])
df[mask]

     A      B  C   D
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, como antes, podemos utilizar numpy para mejorar el rendimiento sin sacrificar prácticamente nada. Usaremos np.in1d

mask = np.in1d(df['A'].values, ['foo'])
df[mask]

     A      B  C   D
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
También incluiré otros conceptos mencionados en otras publicaciones como referencia.
Código a continuación 

Cada columna en esta tabla representa un dataframe de diferente longitud 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ás que los tiempos más rápidos parecen ser compartidos entre mask_with_values y mask_with_in1d

res.T.plot(loglog=True)

enter image description here

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
Mirando el caso especial cuando tenemos un solo no objeto dtypepara todo el marco de datos. Código a continuación 

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 más allá de unos cientos de filas.

spec.T.plot(loglog=True)

enter image description here

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)

76
2017-09-11 22:14



Encuentro que la sintaxis de las respuestas anteriores es redundante y difícil de recordar. Pandas introdujo el query() método en v0.13 y yo 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]: 
          a         b         c
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]: 
          a         b         c
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]: 
          a         b         c
3  0.011763  0.022921  0.244186
8  0.116822  0.364564  0.454607

También puede acceder a variables en el entorno anteponiendo un @.

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

41
2018-02-09 01:36



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]: 
     A      B  C   D
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

11
2017-07-05 16:34



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

10
2018-06-13 11:49



Intenté editar esto, pero no estaba conectado, por lo que no estoy seguro de dónde fue mi edición. Intentaba incorporar una selección múltiple. Entonces creo que una mejor respuesta es:

Para un único valor, el más directo (legible para personas) es probablemente:

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

Para listas de valores, también puede usar:

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)
#      A      B  C   D
# 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

     A      B  C   D
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 sobre los que desea seleccionar, puede ponerlos en una lista y usar 'isin':

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

rendimientos

      A      B  C   D
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

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

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

rendimientos

  A      B  C   D
foo    one  0   0
foo    two  2   4
foo    two  4   8
foo    one  6  12
foo  three  7  14

7
2018-01-25 23:27



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:
   A      B  C   D
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
2018-03-06 06:02