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

Question

我有一个包含 50k 个字符串的数组，称为 products`- 和一个包含大约 2200 万行的数据框，称为 all

我想遍历数组，然后选择包含数组值的数据帧的对应子集：

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

每个查询大约需要 1.5 秒的时间来计算，我的数组中有 50k 个值，这大约需要 20 小时。

你知道计算这个更快的方法吗？

Answer 1

如果你想在产品列表中选择所有带有 id 的行，这应该比 for 循环快得多：

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

Answer 2

为了测试这一点，我生成了我自己版本的这些数据帧（不确定统计属性是否相同，但计时结果似乎与您得到的相似）：

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'])

二分查找法

一种方法是对您的 all 数据框进行排序并使用 searchsorted 将查询作为二分搜索进行 - 对 2200 万行 (n log n) 进行排序的一次性成本很高，但会使查找速度更快 (log n)。这可能是实现您明确目标的最快方法：

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

因此，您似乎可以期望在大约 12 秒内对行进行排序，然后在大约 20 秒内查询 50,000 行，总共需要 32 秒。在我的示例中，我实际上并没有保存结果，但我假设一旦您将索引放入all_products 数据帧（不要称它为all，因为这是内置的Python！），您可以根据需要存储它们。

分组方法

如果all_products 包含来自products 的值全部 或大部分 （就像我的那样），另一种方法（根据我的测试）要快得多, 是将all_products 分组id 并将结果转储到字典中（或任何你想用它做的）：

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

请注意，在这种情况下，它显然比 searchsorted 方法快（尽管速度不快），并且首先不需要对输入进行排序。

请注意，如果您真正想做的是转换这些行或以某种方式修改它们，在这种情况下groupby 绝对是要走的路 - 甚至不要费心转储到字典，而是查看 @ 987654321@ 了解以这种方式使用 Dataframe 的策略。

朴素的方法

为了比较，这里有两种涉及完整搜索的方法：

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