Hay varias formas de seleccionar filas de un marco de datos de pandas:
- Indización booleana (
df[df['col'] == value
])
- Indexación posicional (
df.iloc[...]
)
- Indización de etiquetas (
df.xs(...)
)
df.query(...)
API
A continuación les muestro ejemplos de cada uno, con consejos sobre cuándo usar ciertas técnicas. Supongamos que nuestro criterio es la columna 'A'
=='foo'
(Nota sobre el rendimiento: para cada tipo de base, podemos simplificar las cosas mediante el uso de la API de pandas o podemos aventurarnos fuera de la API, por lo general numpy
, y acelerar las cosas).
Configuración
Lo primero que necesitaremos es identificar una condición que actuará como nuestro criterio para seleccionar filas. Comenzaremos con el caso del OP column_name == some_value
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})
1. indexación booleana
... la indexación booleana requiere encontrar el valor verdadero de la 'A'
columna de cada fila que sea igual 'foo'
, y luego usar esos valores de verdad para identificar qué filas mantener. Por lo general, nos nombramos esta serie, una serie de valores de verdad, mask
. Lo haremos aquí también.
mask = df['A'] == 'foo'
Entonces 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 realizar esta tarea y si el rendimiento o la intuición no son un problema, este debería ser el método elegido. Sin embargo, si el rendimiento es una preocupación, es posible que desee considerar una forma alternativa de crear el mask
.
2. indexación posicional
La indexación posicional ( df.iloc[...]
) 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 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. Indización de etiquetas
La indexación de etiquetas 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)
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. df.query()
API
pd.DataFrame.query
es una forma muy elegante / intuitiva de realizar esta tarea, pero a menudo es más lenta. Sin embargo , si presta atención a los tiempos 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"')
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 cómo creamos nuestro Boolean
mask
.
mask
alternativa 1
Use la numpy
matriz subyacente y renuncie a la sobrecarga de crear otrapd.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, observamos la diferencia en la creación demask
%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 mask
con el numpy
conjunto es ~ 30 veces más rápido. Esto se debe en parte a que la numpy
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 pd.Series
objeto correspondiente .
A continuación, veremos el momento para cortar con uno mask
versus 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 mantiene sobre pruebas más robustas.
mask
alternativa 2
Podríamos haber reconstruido el marco de datos también. Hay una gran advertencia al reconstruir un marco de datos: debe tener cuidado 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 marco de datos es de tipo mixto, que es nuestro ejemplo, cuando obtengamos 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 astype(df.dtypes)
y elimina cualquier ganancia potencial 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 marco de datos 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
La 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 evalúa lo mismo si nuestro conjunto de valores es un conjunto de un valor, a saber 'foo'
. Pero también generaliza para incluir conjuntos de valores más grandes si es necesario. Resulta que esto sigue siendo bastante rápido, aunque es una solución más general. La única pérdida real está 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. Usaremosnp.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
El tiempo
incluiré otros conceptos mencionados en otras publicaciones también como referencia.
Código abajo
Cada columna en esta tabla representa un marco de datos 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 dada 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
ymask_with_in1d
res.T.plot(loglog=True)
Las 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)
Momento especial
Observando el caso especial cuando tenemos un único no objeto dtype
para todo el marco de datos.
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 más allá de unos cientos de filas.
spec.T.plot(loglog=True)
Las 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)