Search value in DataFrame based on multiple columns

Question

Problem: Extract the value of a specific column (in this case 'rating'), for multiple column-value constraints.

Starting with a DataFrame looking as follows

My data is as follows:

    userID  movieID rating
0   196 242 3
1   186 302 3
2   22  377 1

Now, I want to extract the rating for the following case:

userID == 196
movieID == 242

Result should be 3.

I solved it, using the following code: However this isn't very efficient. Does anyone have a better approach?
df.loc[df['userID'] == 196].where(df['movieID'] == 242).dropna()['rating']
Which gives me the rating for the movie with ID 242, and user ID 196.

user_4333359 · Answer

Index lookup is very fast in pandas, so it's best to use that when possible. If users can only rate each movie once a MultiIndex is ideal for this.

df = df.set_index(['userID', 'movieID'])
df.at[(196, 242), 'rating']
#3

Some timings. Once the index is set, lookup is very fast.

%timeit df.at[(userID, movieID), 'rating']
#19.9 µs ± 405 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

df1 = df.reset_index()
%timeit df1.loc[df1['userID'].eq(196) & df1['movieID'].eq(242), 'rating']
#1.2 ms ± 6.98 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

It does take time to set the MultiIndex so that could be costly for a single query. But for many it will quickly pay off, especially with a larger DataFrame. Here's an example of timings where we can use a unique MulitIndex (after the index is set).

import perfplot
import pandas as pd
import numpy as np

perfplot.show(
    setup=lambda n: pd.DataFrame({'userID': range(n),
                                  'movieID': range(n),
                                  'rating': range(n)}).set_index(['userID', 'movieID']), 
    kernels=[
        lambda df: df.at[(4 ,4), 'rating'],
        lambda df: df.loc[(df.index.get_level_values('userID') == 4) 
                          & (df.index.get_level_values('movieID') == 4), 'rating']
    ],
    labels=["MultiIndex", "Boolean Slice"],
    n_range=[2 ** k for k in range(5, 25)],
    equality_check=np.allclose,  
    xlabel="len(df)"
)

Search value in DataFrame based on multiple columns

1 Answers