如何加速 pandas 为 numba 引擎应用 lambda 函数

Question

对于加速 pandas 的案例，我了解了 numba 引擎，可以显着加速。

在最近的情况下，我喜欢使用带有 lambda 的 argsort 来获取任意索引，但它似乎很慢。为什么 lambda 会减慢代码速度？如何编写适当的应用函数以不进一步损害速度？？我什至将 lambda 函数包装在 nb.njit 中，但仍然看不到更多的加速。

从逻辑上讲 np.argmax 和 np.argsort [-1] 是相同的，但是例如可以使用 np.argsort 来获取中位数。所以我想确认是否有更好的方法来写lambda x: np.argsort(x)[5]，例如

import pandas as pd
import numpy as np
import numba as nb
import timeit
import sys


def f1():
    a = pd.DataFrame(range(10000000))
    return a.rolling(10).apply(np.argmax, engine='numba', raw=True)

nb.njit((nb.int64)(nb.float64[:]))
def f(x):
    return np.argsort(x)[5]

def f2():
    a = pd.DataFrame(range(10000000))
    return a.rolling(10).apply(f, engine='numba', raw=True)
if __name__ == '__main__':
    print(timeit.timeit(lambda: f2(), number=10) / 10)

(base) xxx:~$ python test.py f1
1.4400632409029641
(base) xxx:~$ python test.py f2
5.061740894208197

Answer 1

这些是我使用timeit得到的结果

In [1]: %timeit a.rolling(10).apply(np.argmax, engine='numba', raw=True)
16.9 ms ± 477 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [2]: %timeit a.rolling(10).apply(np.argmax, raw=True)
176 ms ± 3.07 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [3]: %timeit a.rolling(10).apply(l1, engine='numba', raw=True)
79.8 ms ± 1.93 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [4]: %timeit a.rolling(10).apply(l1, raw=True)
116 ms ± 1.65 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

请注意，我使用 a = pd.DataFrame(range(100_000)) 进行了此操作。