pandas 提高大型 Dataframe 的df.rolling(...).apply(...)的性能

vohkndzv  于 2022-11-20  发布在  其他
关注(0)|答案(1)|浏览(423)

此代码的执行时间太长。

df.rolling(window=255).apply(myFunc)

我的 Dataframe 形状是(500,10000)。

0         1 ... 9999
2021-11-01  0.011111  0.054242 
2021-11-04  0.025244  0.003653 
2021-11-05  0.524521  0.099521 
2021-11-06  0.054241  0.138321 
...

我用最后255个日期值对每个日期进行计算。myFunc看起来像:

def myFunc(x):
   coefs = ...
   return np.sqrt(np.sum(x ** 2 * coefs))

我尝试使用更快的速度,但效果相同:

import swifter
df.swifter.rolling(window=255).apply(myFunc)

我也尝试过Dask,但我想我并没有很好地理解它,因为性能并没有好到哪里去:

import dask.dataframe as dd
ddf = dd.from_pandas(df)
ddf = ddf.rolling(window=255).apply(myFunc, raw=False)
ddf.execute()

我没有设法用分区来并行执行。我如何使用dask来提高性能?我使用的是Windows。

p5cysglq

p5cysglq1#

使用numpy + numba可以非常高效地完成此操作。
快速MRE:

import numpy as np, pandas as pd, numba

df = pd.DataFrame(
    np.random.random(size=(500, 10000)),
    index=pd.date_range("2021-11-01", freq="D", periods=500)
)

coefs = np.random.random(size=255)

使用numba.njit(parallel=True)numba.prange,使用纯numpy运算和简单循环编写函数:

@numba.njit(parallel=True)
def numba_func(values, coefficients):
    # define result array: size of original, minus length of
    # coefficients, + 1
    result_tmp = np.zeros(
        shape=(values.shape[0] - len(coefficients) + 1, values.shape[1]),
        dtype=values.dtype,
    )

    result_final = np.empty_like(result_tmp)

    # nested for loops are your friend with numba!
    # (you must unlearn what you have learned)
    for j in numba.prange(values.shape[1]):
        for i in range(values.shape[0] - len(coefficients) + 1):
            for k in range(len(coefficients)):
                result_tmp[i, j] += values[i + k, j] ** 2 * coefficients[k]

        result_final[:, j] = np.sqrt(result_tmp[:, j])

    return result_final

此命令运行非常快:

In [5]: %%time
   ...: result = pd.DataFrame(
   ...:     numba_func(df.values, coefs),
   ...:     index=df.index[len(coefs) - 1:],
   ...: )
   ...:
   ...:
CPU times: user 1.69 s, sys: 40.9 ms, total: 1.73 s
Wall time: 844 ms

注意事项:我是dask的忠实粉丝。但是dask性能的第一条规则是don't use dask。如果它足够小,可以轻松地放入内存中,那么通常可以通过调整panda或numpy操作并利用cython、numba等的加速来获得最佳性能。一旦问题大到可以转移到dask,这些相同的调整规则也适用于在dask块/分区上执行的操作。我也是!

相关问题