提高 df.rolling(...).apply(...) 對大型數據幀的性能

Question

此代碼的執行時間太長。

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

我的數據框形狀是 (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 個日期值對每個日期進行計算。 我的函數看起來像：

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

我嘗試使用 swifter 但性能是一樣的：

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。

Answer 1

這可以使用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)

使用純 numpy 操作和簡單循環編寫 function，利用numba.njit(parallel=True)和numba.prange ：

@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 性能的第一條規則是 不要使用 dask 。 如果它足夠小，可以輕松放入 memory，您通常會通過調整 pandas 或 numpy 操作並利用 cython、numba 等的加速來獲得最佳性能。一旦問題大到足以轉移到 dask，這些相同的調整規則也適用於您在 dask 塊/分區上執行的操作！

Answer 2

首先，由於您使用的是numpy函數，因此指定參數raw=True 。 玩具示例：

import pandas as pd
import numpy as np

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

df = pd.DataFrame(np.random.random((500, 10000)))

%%time
res = df.rolling(250).apply(foo)

Wall time: 359.3 s

# with raw=True
%%time
res = df.rolling(250).apply(foo, raw=True)

Wall time: 15.2 s

您還可以使用parallel-pandas庫輕松並行化計算。 只有兩行額外的代碼！

# pip install parallel-pandas
import pandas as pd
import numpy as np
from parallel_pandas import ParallelPandas

#initialize parallel-pandas
ParallelPandas.initialize(n_cpu=8, disable_pr_bar=True)

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

df = pd.DataFrame(np.random.random((500, 1000)))

# p_apply - is parallel analogue of apply method
%%time
res = df.rolling(250).p_apply(foo, raw=True, executor='processes')

Wall time: 2.2 s

隨着engine='numba'

%%time
res = df.rolling(250).p_apply(foo, raw=True, executor='processes', engine='numba')

Wall time: 1.2 s

總加速比為359/1.2 ~ 300 ！