Fastest ways to filter for values in pandas dataframe

Viewed 796

A similar dataframe can be created:

import pandas as pd

df = pd.DataFrame()
df["nodes"] = list(range(1, 11))
df["x"] = [1,4,9,12,27,87,99,121,156,234]
df["y"] = [3,5,6,1,8,9,2,1,0,-1]
df["z"] = [2,3,4,2,1,5,9,99,78,1]

df.set_index("nodes", inplace=True)

So the dataframe looks like this:

        x        y        z
nodes
1       1        3        2
2       4        5        3
3       9        6        4
4       12       1        2
5       27       8        1
6       87       9        5
7       99       2        9
8       121      1        99
9       156      0        78
10      234      -1       1

My first try for searching e.g. all nodes containing number 1 is:

>>> df[(df == 1).any(axis=1)].index.values
[1 4 5 8 10]

As i have to do this for many numbers and my real dataframe is much bigger than this one, i'm searching for a very fast way to do this.

1 Answers

Just tried something that may be enlightening

import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.randint(0,100,size=(10000, 4)), columns=list('ABCD'))
df.set_index("A", inplace=True)
df_no_index = df.reset_index()

So set up a dataframe with ints right the way through. This is not the same as yours but it will suffice.

Then I ran four tests

%timeit df[(df == 1).any(axis=1)].index.values
%timeit df[(df['B'] == 1) | (df['C']==1)| (df['D']==1)].index.values
%timeit df_no_index[(df_no_index == 1).any(axis=1)].A.values
%timeit df_no_index[(df_no_index['B'] == 1) | (df_no_index['C']==1)| (df_no_index['D']==1)].A.values

The results I got were,

940 µs ± 41.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.47 ms ± 7.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.08 ms ± 14.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.55 ms ± 51.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Which showed that the initial method that you took, with index seems to be the fastest of these approaches. Removing the index does not improve the speed with a moderately sized dataframe

Related