Efficiently filter DataFrame by looking for NumPy array match in row

user_8053020 0
Viewed 513

Given

df = pd.DataFrame({'x': [np.array(['1', '2.3']), np.array(['30', '99'])]},
                  index=[pd.date_range('2020-01-01', '2020-01-02', freq='D')])

I would like to filter for np.array(['1', '2.3']). I can do

df[df['x'].apply(lambda x: np.array_equal(x, np.array(['1', '2.3'])))]

but is this the fastest way to do it?

EDIT: Let's assume that all the elements inside the numpy array are strings, even though it's not good practice!

DataFrame length can go to 500k rows and the number of values in each numpy array can go to 10.

pandas numpy
2 Answers

You can rely on list comprehension for performance:

df[np.array([np.array_equal(x,np.array([1, 2.3])) for x in df['x'].values])]

Performance via timeit(on my system currently using 4gb ram) :

%timeit -n 2000 df[np.array([np.array_equal(x,np.array([1, 2.3])) for x in df['x'].values])]
#output:
425 µs ± 10.8 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)

%timeit -n 2000 df[df['x'].apply(lambda x: np.array_equal(x, np.array([1, 2.3])))]
#output:
875 µs ± 28.6 µs per loop (mean ± std. dev. of 7 runs, 2000 loops each)
user_142898920

My suggestion would be to do the following:

import numpy as np
mat = np.stack([np.array(["a","b","c"]),np.array(["d","e","f"])])

In reality this would be the actual data from the cols of your dataframe. Make sure that these are a single numpy array.

Then do:

 matching_rows = (np.array(["a","b","c"]) == mat).all(axis=1)

Which outputs you an array of bools indicating where the matches are located. So you can then filter your rows like this:

df[matching_rows]
user_43108630
Related