Pandas 将 .apply(lambda x: 替换为快速解决方案，例如使用 numpy 数组

Question

我正在尝试加快我用来处理具有数百个因素的数百万行的排名函数。我在下面提供了一个示例 MCVE：

to_rank = ['var_1', 'var_2', 'var_3']
df = pd.DataFrame({'var_1' : np.random.randn(1000), 'var_2' : np.random.randn(1000), 'var_3' : np.random.randn(1000)})
df['date_id'] = np.random.choice(range(2001, 2012), df.shape[0])
df['category'] = ','.join(chr(random.randrange(97, 97 + 4 + 1)).upper() for x in range(1,df.shape[0]+1)).split(',')

我的排名代码如下：

import pandas as pd
import numpy as np
import bottleneck as bn

%timeit ranked = df[['date_id', 'category'] + to_rank].groupby(['date_id', 'category']).apply(lambda x: x[to_rank].apply(lambda x: bn.nanrankdata(x) * 100 / len(x) - 1))

10 loops, best of 3: 106 ms per loop

根据我的数据，这大约需要 30 到 40 秒。我收集到 .apply(lambda x: 有很大的开销，包括循环、dtype 检测和形状分析，我使用它两次循环多索引，这可能效率低下。我读过一个可以通过使用 Series/numpy 数组对其进行矢量化（例如https://tomaugspurger.github.io/modern-4-performance.html，但我自己正在努力实现这一点；事实上，关于在多索引上应用函数的大多数类似问题似乎都使用 .apply(lambda x: 所以我怀疑其他人可以也可以从加速他们的代码中受益。

Answer 1

您可以定义一个函数并使用transform，尽管所花费的时间并没有那么好（只有两倍快）：

def nanrankdata_len(x):
    return bn.nanrankdata(x)*100/len(x) - 1

%timeit ranked = df.groupby(['date_id','category']).transform(nanrankdata_len)
#-> 10 loops, best of 3: 55.5 ms per loop