web-dev-qa-db-fra.com

Sélectionnez efficacement les lignes qui correspondent à une ou plusieurs valeurs dans Pandas DataFrame

Problème

Étant donné les données dans un Pandas DataFrame comme le suivant:

Name     Amount
---------------
Alice       100
Bob          50
Charlie     200
Alice        30
Charlie      10

Je souhaite sélectionner toutes les lignes où Name est l'une des nombreuses valeurs d'une collection {Alice, Bob}

Name     Amount
---------------
Alice       100
Bob          50
Alice        30

Question

Quel est un moyen efficace de le faire dans Pandas?

Options telles que je les vois

  1. Parcourez les lignes, manipulant la logique avec Python

  2. Sélectionnez et fusionnez de nombreuses instructions comme les suivantes

    merge(df[df.name = specific_name] for specific_name in names) # something like this

  3. Effectuer une sorte de jointure

Quels sont les compromis de performance ici? Quand une solution est-elle meilleure que les autres? Quelles solutions me manquent?

Bien que l'exemple ci-dessus utilise des chaînes, mon travail réel utilise des correspondances sur 10 à 100 entiers sur des millions de lignes et des opérations NumPy rapides peuvent donc être pertinentes.

pythonpandas
36
MRocklin

Vous pouvez utiliser la méthode isin Series:

In [11]: df['Name'].isin(['Alice', 'Bob'])
Out[11]: 
0     True
1     True
2    False
3     True
4    False
Name: Name, dtype: bool

In [12]: df[df.Name.isin(['Alice', 'Bob'])]
Out[12]: 
    Name  Amount
0  Alice     100
1    Bob      50
3  Alice      30
56
Andy Hayden

Étant donné que dans votre cas d'utilisation réel, les valeurs dans df['Name'] Sont ints, vous pourrez peut-être générer le masque booléen plus rapidement en utilisant l'indexation NumPy au lieu de Series.isin.

idx = np.zeros(N, dtype='bool')
idx[names] = True
df[idx[df['Name'].values]]

Par exemple, compte tenu de cette configuration:

import pandas as pd
import numpy as np

N = 100000
df = pd.DataFrame(np.random.randint(N, size=(10**6, 2)), columns=['Name', 'Amount'])
names = np.random.choice(np.arange(N), size=100, replace=False)

In [81]: %timeit idx = np.zeros(N, dtype='bool'); idx[names] = True; df[idx[df['Name'].values]]
100 loops, best of 3: 9.88 ms per loop

In [82]: %timeit df[df.Name.isin(names)]
10 loops, best of 3: 107 ms per loop

In [83]: 107/9.88
Out[83]: 10.82995951417004

N est (essentiellement) la valeur maximale que df['Names'] peut atteindre. Si N est plus petit, l'avantage de vitesse n'est pas aussi important. Avec N = 200,

In [93]: %timeit idx = np.zeros(N, dtype='bool'); idx[names] = True; df[idx[df['Name'].values]]
10 loops, best of 3: 62.6 ms per loop

In [94]: %timeit df[df.Name.isin(names)]
10 loops, best of 3: 178 ms per loop

In [95]: 178/62.6
Out[95]: 2.8434504792332267

Attention: comme indiqué ci-dessus, il semble y avoir un avantage de vitesse, en particulier lorsque N devient grand. Cependant, si N est trop grand, la formation de idx = np.zeros(N, dtype='bool') peut ne pas être possible.

Verification sanitaire:

expected = df[df.Name.isin(names)]
idx = np.zeros(N, dtype='bool')
idx[names] = True
result = df[idx[df['Name'].values]]
assert expected.equals(result)
5
unutbu