Existem várias maneiras de selecionar linhas de um quadro de dados do pandas:
- Indexação booleana (
df[df['col'] == value])
- Indexação posicional (
df.iloc[...])
- Indexação de etiquetas (
df.xs(...))
df.query(...) API
Abaixo, mostro exemplos de cada um, com conselhos sobre quando usar determinadas técnicas. Suponha que nosso critério seja a coluna 'A'=='foo'
(Observação sobre desempenho: para cada tipo de base, podemos simplificar as coisas usando a API do pandas ou podemos nos aventurar fora da API, geralmente numpy, e acelerar as coisas.)
Configuração
A primeira coisa que precisamos é identificar uma condição que atuará como nosso critério para selecionar linhas. Começaremos com o caso do OP column_name == some_valuee incluiremos outros casos de uso comuns.
Empréstimos 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. Indexação booleana
... A indexação booleana exige que o valor verdadeiro da 'A'coluna de cada linha seja igual a e 'foo', em seguida, use esses valores verdadeiros para identificar quais linhas manter. Normalmente, nomearíamos essa série, uma matriz de valores de verdade mask,. Faremos isso aqui também.
mask = df['A'] == 'foo'
Podemos então usar essa máscara para cortar ou indexar o quadro de dados
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
Essa é uma das maneiras mais simples de realizar essa tarefa e, se o desempenho ou a intuição não forem um problema, esse deve ser o método escolhido. No entanto, se o desempenho for uma preocupação, convém considerar uma maneira alternativa de criar o mask.
2. Indexação posicional
A indexação posicional ( df.iloc[...]) tem seus casos de uso, mas este não é um deles. Para identificar onde dividir, primeiro precisamos executar a mesma análise booleana que fizemos acima. Isso nos deixa executando uma etapa extra para realizar a mesma tarefa.
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. Indexação de etiquetas
A indexação de etiquetas pode ser muito útil, mas, neste caso, estamos trabalhando novamente sem nenhum benefício
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é uma maneira muito elegante / intuitiva de executar essa tarefa, mas geralmente é mais lenta. No entanto , se você prestar atenção aos horários abaixo, para dados grandes, a consulta é muito eficiente. Mais do que a abordagem padrão e de magnitude semelhante à minha melhor sugestão.
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
Minha preferência é usar o Boolean mask
As melhorias reais podem ser feitas modificando a forma como criamos o nosso Boolean mask.
maskalternativa 1
Use a numpymatriz subjacente e renuncie à sobrecarga de criar outrapd.Series
mask = df['A'].values == 'foo'
Mostrarei testes de tempo mais completos no final, mas dê uma olhada nos ganhos de desempenho que obtemos usando o quadro de dados de amostra. Primeiro, analisamos a diferença na criação domask
%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)
Avaliar o maskcom a numpymatriz é ~ 30 vezes mais rápido. Isso se deve em parte à numpyavaliação frequentemente mais rápida. Também se deve em parte à falta de sobrecarga necessária para criar um índice e um pd.Seriesobjeto correspondente .
A seguir, veremos o momento para fatiar com um maskversus o outro.
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)
Os ganhos de desempenho não são tão pronunciados. Vamos ver se isso se sustenta em testes mais robustos.
maskalternativa 2
Poderíamos ter reconstruído o quadro de dados também. Há uma grande ressalva ao reconstruir um quadro de dados - você deve cuidar disso dtypesao fazer isso!
Em vez de df[mask]fazermos isso
pd.DataFrame(df.values[mask], df.index[mask], df.columns).astype(df.dtypes)
Se o quadro de dados for do tipo misto, como é o nosso exemplo, quando obtermos df.valuesa matriz resultante for dtype objecte, consequentemente, todas as colunas do novo quadro de dados serão de dtype object. Exigindo, assim, astype(df.dtypes)e eliminando possíveis ganhos de desempenho.
%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)
No entanto, se o quadro de dados não for do tipo misto, essa é uma maneira muito útil de fazê-lo.
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)
Reduzimos o tempo pela metade.
maskalternativa 3 O
@unutbu também nos mostra como usar pd.Series.isinpara contabilizar cada elemento de df['A']estar em um conjunto de valores. Isso avalia o mesmo se nosso conjunto de valores é um conjunto de um valor, a saber 'foo'. Mas também generaliza para incluir conjuntos maiores de valores, se necessário. Acontece que isso ainda é muito rápido, mesmo que seja uma solução mais geral. A única perda real está na intuição para aqueles que não estão familiarizados com o conceito.
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
No entanto, como antes, podemos utilizar numpy para melhorar o desempenho sem sacrificar praticamente nada. Vamos usarnp.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
Cronometragem
Vou incluir outros conceitos mencionados em outras postagens e também para referência.
Código Abaixo
Cada coluna nesta tabela representa um quadro de dados de comprimento diferente, sobre o qual testamos cada função. Cada coluna mostra o tempo relativo gasto, com a função mais rápida com um í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
Você notará que os tempos mais rápidos parecem ser compartilhados entre mask_with_valuesemask_with_in1d
res.T.plot(loglog=True)

Funções
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]
Teste
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
Analisando o caso especial em que temos um único não-objeto dtypepara todo o quadro de dados.
Código Abaixo
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
Acontece que a reconstrução não vale a pena além de algumas centenas de linhas.
spec.T.plot(loglog=True)

Funções
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
)
Teste
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)