Pandas Query Optimization On Multiple Columns

Question

My training dataset is around 5 MB and test dataset is of the same size. I have run some simulations over the whole dataset couple of times. However, in a latter solution, I ran queries on two columns (say A and B). Now, in the calculation, for each row in the test dataset, I have to get the result of the following query. This time, the program takes forever. I can optimize this on my own writing some extra code (sorting will help in my case and in fact I can skip pandas altogether if I want to), but I wanted to check if there is a better built-in solution should I ever require that (also, I don't want to if I don't have to). The query is like this:

def f(A, B, x, y):
    data = df.loc[df[A].isin([x]) & df.B.isin([y])]
    return len(data)

Here, df is the DataFrame (training dataset). For each row in test dataset, data find all rows where column A has value x and B has value y. As far as I know, isin is slightly faster, so I used it. However, I think this is why it does not run within a second like it did before because each time I run this query, it needs a lot of time. What I have to do next is something like this this.

for i in range(0,l):
    ret = f(A, B, test[A].values[i],test[B].values[i])

Clearly, in every loop, it will calculate the result of the query. Other than this, rest of the functions in my code are mathematical that do not need much calculation.

Let me know if the problem is not clear since I did not use direct codes (replaced it with dummy variables and other lines do not have anything that has high complexity).

Answer 1

Short answer

There are far faster ways vectorized ways to do this than brute-force loops. One particular option while remaining Pandas-level would be

(tra_df.groupby(tra_df.columns.tolist())
       .size()
       .reindex(tst_df.values.T.tolist(), fill_value=0)

This should offer you an enormous performance boost, which could be further improved with a NumPy vectorized solution, depending on what you're satisfied with.

---

Long answer

Suppose that your training and test data looked something like this.

import pandas as pd; import numpy as np; from numpy.random import RandomState

rnd_st = RandomState(8675309)

tra_df = pd.DataFrame(dict(A=rnd_st.randint(0, 10, 10**3), 
                           B=rnd_st.randint(0, 10, 10**3)))
tst_df = pd.DataFrame(dict(A=rnd_st.randint(0, 10, 10**3), 
                           B=rnd_st.randint(0, 10, 10**3)))

First of all, you aren't really using Pandas correctly with

def f(A, B, x, y):
    data = df.loc[df[A].isin([x]) & df.B.isin([y])]
    return len(data)

isin isn't intended to match against single values like this, but rather a list or set of values when you don't have another option. In your case, you can achieve the same boolean index directly with

(df.A == x) & (df.B == y)

Furthermore, there's no point indexing your DataFrame again with df.loc only to take the length of it... you might as well just use sum directly on the boolean vector to count the matches. Putting this together and using apply over the test DataFrame, you can achieve your same output with

tst_df.apply(lambda x: ((tra_df.A == x.A) & (tra_df.B == x.B)).sum(), axis=1)

But this is still very slow - wrapping the above approach in a function and comparing it to your original approach on the sample data

%timeit your_way(tra_df, test_df)
1 loops, best of 3: 819 ms per loop
%timeit direct_er_way(tra_df, test_df)
1 loops, best of 3: 673 ms per loop

Much faster would be to use groupby and then reindex, as instead of brute-force looping this offers a vectorized solution where we are effectively hashing the counts.

(tra_df.groupby(tra_df.columns.tolist())
       .size()
       .reindex(tst_df.values.T.tolist(), fill_value=0)

Benchmarking,

 %timeit groupby_way()
100 loops, best of 3: 3.04 ms per loop

I'm not that well-versed in NumPy, but I can safely assume that were this function still not fast enough to meet your needs then a NumPy vectorized solution avoiding some of the overhead would be the next step.