Improve performance of Pandas subsets/filters with multiple conditions

Question

In my Python script I have a Pandas DataFrame with about 5.2 million rows and 26 columns.

The current analysis I am running tests 1.5 million different subsets/filters against this DataFrame, using multiple conditions (combinations of the columns) to count the total number of occurrences. I've tried multiple methods and improved the speed a bit from my original code, but I believe this could be done faster.

Heres the code with a dummy DataFrame (5.2M rows and 26 columns A-Z) and one example of the subset method I initially used, with multiple conditions:

import pandas as pd
import numpy as np
import time

np.random.seed(101)
df = pd.DataFrame(np.random.randint(0,11,size=(5200000, 26)), columns=list('ABCDEFGHIJKLMNOPQRSTUVWXYZ'))

cond1 = (df.A > 5)
cond2 = (df.B.shift(1).between(8,10,inclusive=True))
cond3 = (df.C.shift(1) < 5)
cond4 = (df.D.shift(1) == 1)
cond5 = (df.E.shift(1).between(2,5,inclusive=True))

start = time.time()
subsetLength = len(df.loc[cond1 & cond2 & cond3 & cond4 & cond5])
stop = time.time()

print("Length:",subsetLength)
print("Time:", stop - start)

The time to retrieve the length of this subset is around 0.030s on my laptop (Macbook Pro 2020 Intel) for each subset/filter:

Length: 9781
Time: 0.029639244079589844

Note: I am using .shift(1) to check previous rows for the columns B - E and for column A this is not the case.

Now 0.030s might seem quite fast, but since I am running 1.5 million of those subsets, the script currently took about 12-13 hours. Therefor, every speed increase I can get at an individual subset/filter can be a huge time saver.

Instead of .loc[] I have also tried .shape[0] and .sum() like this:

.shape[0]

subsetLength = df[cond1 & cond2 & cond3 & cond4 & cond5].shape[0]

This seems to be a bit faster but still was in the range of 0.027s - 0.028s.

.sum()

subsetLength = (cond1 & cond2 & cond3 & cond4 & cond5).sum()

This was around 0.021s, already quite a nice improvement compared to the .loc[] which could save around 4 hours. But still the total would take a lot of hours (around 8 or 9).

np.where()

Finally I tried to use np.where(), for this I had to change the conditions a bit:

cond1 = (df['A'] > 5)
cond2 = (df['B'].shift(1).between(8,10,inclusive=True))
cond3 = (df['C'].shift(1) < 5)
cond4 = (df['D'].shift(1) == 1)
cond5 = (df['E'].shift(1).between(2,5,inclusive=True))

subsetLength = len(np.where(cond1 & cond2 & cond3 & cond4 & cond5)[0])

This took around 0.0195s and was the fastest method I got thus far.

Preferably I am looking for performance in the 0.002s or faster range so the script would take around 1 hour to finish, but I can't tell if this is possible. Any advice on how the speed of this subsetting/filtering a big DataFrame for counts can be improved? Am I perhaps using slow ways to do counts for subsets with multiple conditions? Or is this kind of analysis just heavy and is the time it takes realistic.

user_12939557 · Answer

One solution is to perform the logical operation in-place to avoid the creation of many temporary arrays:

tmp = np.empty(cond1.shape, dtype=np.bool8)
np.logical_and(cond1, cond2, out=tmp)
np.logical_and(tmp, cond3, out=tmp)
np.logical_and(tmp, cond4, out=tmp)
np.logical_and(tmp, cond5, out=tmp)
subsetLength = tmp.sum()

This solution is 2.2 times faster on my machine.

The above solution is not very efficient as it still read and write big arrays in memory. However, I doubt it is possible to do something much faster only with Numpy (since logical_and apparently do not provide a way to compute more than 2 elements at a time).

Another, solution is to use Numba to avoid the read/write of the temporary array and even run the code in parallel. Here is how:

import numba as nb

@nb.njit('int64(bool_[::1],bool_[::1],bool_[::1],bool_[::1],bool_[::1])', parallel=True)
def computeSubsetLength(cond1, cond2, cond3, cond4, cond5):
    n = len(cond1)
    assert len(cond2) == n and len(cond3) == n and len(cond4) == n and len(cond5) == n
    subsetLength = 0
    for i in nb.prange(n):
        subsetLength += cond1[i] & cond2[i] & cond3[i] & cond4[i] & cond5[i]
    return subsetLength

subsetLength = computeSubsetLength(cond1.to_numpy(), cond2.to_numpy(), cond3.to_numpy(), cond4.to_numpy(), cond5.to_numpy())

This new version is 31 times faster. It is close to optimal on my machine since it just takes roughly the time to read the arrays in memory (is takes 0.000624 s while the time to only read the input data from memory is 0.000484 s).

Improve performance of Pandas subsets/filters with multiple conditions

1 Answers