parallelize/speed up pandas groupby 变换的高效方法

Question

我正在尝试加速一些用于计算 pandas 时间索引数据帧滞后的代码。数据框包含由 ID 列标识的约 200k 时间序列。我尝试了 dask 但没有任何改进（比单独 pandas 花费的时间更长）。

这是一个生成具有可比较大小的虚拟数据框的工作示例：

import itertools as it
import numpy as np
import pandas as pd
np.random.seed(1)

#Series for ID
ID_data = pd.Series(np.arange(0,200000), name='ID')

#Array of data - create pandas dataframe with datetime index
value_data = np.random.rand(52,1)
tidx = pd.date_range('2019-01-01', periods=len(value_data), freq='D')
#Cross join with ID to create test dataframe
df = pd.DataFrame(value_data, columns=['value'], index=tidx).reset_index().merge(ID_data,how="cross").set_index('index')

现在，我想为每个时间序列（由 ID 列标识）计算值列的滞后（在此示例中为 1 天滞后）：

%%time
df["value_lag1"] = df.groupby(['ID'])["value"].transform(lambda x: x.shift(1))

此代码需要 30 秒才能执行。你知道有什么有效的方法可以加快速度吗？

谢谢最好的问候

Answer 1

您可以通过 .transform 在不使用 lambda 函数的情况下缩短执行时间。 GroupBy对象直接使用DataFrameGroupBy.shift()函数即可，如下：

df["value_lag1"] = df.groupby(['ID'])["value"].shift(1)

原始版本与此版本在我机器上的执行时间是36.6s vs 0.715s。提高 51 倍。

使用 .transform 中的 lambda 函数，您没有使用内置矢量化 Pandas 操作，而是使用缓慢的非优化代码。通过直接使用 DataFrameGroupBy.shift() 函数，您的代码变得矢量化并且运行更快。

结果比较

2个新旧代码生成的列名：

df["value_lag1"] = df.groupby(['ID'])["value"].transform(lambda x: x.shift(1))

df["value_lag2"] = df.groupby(['ID'])["value"].shift(1)

df["value_lag1"].compare(df["value_lag2"])

# No difference shown by the compare function:

        self    other
index

性能比较

%%timeit
df["value_lag1"] = df.groupby(['ID'])["value"].transform(lambda x: x.shift(1))

36.6 s ± 768 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
df["value_lag1"] = df.groupby(['ID'])["value"].shift(1)

715 ms ± 64.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

36.6 秒与 0.715 秒：提高 51 倍

parallelize/speed up pandas groupby 变换的高效方法

Efficient way to parallelize/speed up pandas groupby transform

python

parallel-processing

time-series

dataframe

pandas

结果比较

性能比较