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_value
e 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
.
mask
alternativa 1
Use a numpy
matriz 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 mask
com a numpy
matriz é ~ 30 vezes mais rápido. Isso se deve em parte à numpy
avaliação frequentemente mais rápida. Também se deve em parte à falta de sobrecarga necessária para criar um índice e um pd.Series
objeto correspondente .
A seguir, veremos o momento para fatiar com um mask
versus 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.
mask
alternativa 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 dtypes
ao 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.values
a matriz resultante for dtype
object
e, 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.
mask
alternativa 3 O
@unutbu também nos mostra como usar pd.Series.isin
para 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_values
emask_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 dtype
para 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)