通过Pandas Dataframe进行迭代和索引的最快方法

Question

I have an array of 50k strings called products`- and a dataframe of about 22 million rows called all

I want to iterate through the array and then select a corresponding subset of the dataframe that contains the array value:

for i in products:
 all.query('id == i')

Each query takes about 1.5s to compute, with 50k values in my array that will take me about 20 hours.

Do you know any faster way to compute this ?

Answer 1

If you want to select all rows with ids in the products list, this should be much faster than a for loop:

import numpy as np    
df[np.in1d(df.id,products)]

Answer 2

In order to test this, I generated my own version of these dataframes (not sure if the statistical properties are the same, but the timing results seem similar to what you're getting):

import pandas as pd
import numpy as np

import uuid

products = pd.Series([uuid.uuid4().hex for i in range(50000)])
all_products = pd.DataFrame(np.random.choice(products,
                                             size=(int(22e6),), replace=True),
                            columns=['id'])

Binary search method

One way to do this is to sort your all dataframe and use searchsorted to do the queries as binary searches - which has a one-time heavy cost sorting the 22M rows ( n log n ), but makes the lookups much faster ( log n ). This may be the fastest way to achieve your explicitly stated goal:

import timeit
s = timeit.default_timer()
all_products_sorted = all_products.sort_values(by='id')
e = timeit.default_timer()
print('Time to sort: {:0.5f}'.format((e - s) / N))
# Time to sort: 11.27207

N = 1000
s = timeit.default_timer()
for _, i in zip(range(N), products):
    start = all_products_sorted['id'].searchsorted(i, side='left')
    end = all_products_sorted['id'].searchsorted(i, side='right')
    x = all_products_sorted['id'].iloc[start[0]:end[0]]
e = timeit.default_timer()

print('{:0.5f}s per query'.format((e - s) / N))
# 0.00038s per query

So it seems that you can expect to sort the rows in around 12s, and then query the 50,000 rows in another ~20s, for a total of 32s. In my example I don't actually save the results, but I assume once you have the indices into the all_products dataframe (don't call it all because that's a Python builtin!), you can store them as desired.

Groupby Method

Another method, which (according to my test), is considerably faster if all_products consists entirely or mostly of values from products (as mine does), is to group all_products by id and dump the result into a dictionary (or whatever you want to do with it):

s = timeit.default_timer()
x_dict = {k: v for k, v in all_products.groupby('id')}
e = timeit.default_timer()
print('{:0.5f}s per query'.format((e - s) / len(products)))
# 0.00032s per query

Note that in this case it is apparently faster than the searchsorted method (though not considerably), and doesn't require the input to be sorted in the first place.

Note that if what you actually want to do is transform these rows or modify them in some way, in this case groupby is definitely the way to go - don't even bother dumping to a dictionary, instead see the split-apply-combine page for strategies on working with Dataframes in this way.

Naive methods

For comparison, here are two approaches that involve full searches:

import timeit
N = 5
s = timeit.default_timer()
for _, i in zip(range(N), products):
    x = all_products.query('id == "{}"'.format(i))
e = timeit.default_timer()

print('{:0.5f}s per query'.format((e - s) / N))  # 1.60075s per query


s = timeit.default_timer()
for _, i in zip(range(N), products):
    x = all_products[all_products['id'] == i]
e = timeit.default_timer()

print('{:0.5f}s per query'.format((e - s) / N))  # 3.00135s per query