通过聚合加快熊猫数据帧分组的最佳方法

时间:2020-10-22 17:52:47

标签: python pandas dataframe optimization multiprocessing

我有一个像这样的大数据框price_d

+---------------------------------------------------+
| date          monthEndDate  stock  volume  logRet |
+---------------------------------------------------+
| 1990-01-01    1990-01-31    A      1       NA     |
| 1990-01-02    1990-01-31    A      2       0.2    |
| 1990-02-01    1990-02-28    A      3       0.3    |
| 1990-02-02    1990-02-28    A      4       0.4    |
| ...           ...                                 |
| 1990-01-01    1990-01-31    B      1       NA     |
| 1990-01-02    1990-01-31    B      2       0.08   |
| ...           ...                                 |
| 1990-02-01    1990-02-28    B      0       0.3    |
| 1990-02-02    1990-02-28    B      3       0.4    |
| ...           ...                                 |
+---------------------------------------------------+

此数据帧的长度将为百万,其中monthEndDate中有数百个不同的值,而stock中有数千个不同的值。

我使用三个自定义函数对volume和logRet进行了分组汇总:

def varLogRet(_s):
    return pd.Series({'varLogRet': np.var(_s.iloc[_s.to_numpy().nonzero()])})

def TotRet1M(_s):
    return pd.Series({'TotRet1M': np.exp(np.sum(_s))-1})


def avgVolume(_s):
    return pd.Series({'avgVolume': np.mean(_s.iloc[_s.to_numpy().nonzero()])})

return_m = price_d.groupby(['monthEndDate', 'tradingItemId']).agg({'logRet': [varLogRet, TotRet1M],
                                                                       'volume': avgVolume})

groupby聚合将花费数分钟的时间。以我为例,加快此过程的最佳方法是什么,多处理是否可行?

1 个答案:

答案 0 :(得分:0)

当内置有熊猫并且可以直接使用可能的优化功能时,您真的不需要.agg。默认情况下,NaN被忽略。只需分别计算所需的列,然后再使用它们即可。

基准测试:在运行64位debian 10的我的普通Core i5-8250U(4C8T)笔记本电脑上,完成了800万行数据,花费了不到3s的时间。

# make a dataset of 8 million rows
df = pd.read_clipboard(sep=r"\s{2,}")
df2 = df.loc[df.index.repeat(1000000)].reset_index(drop=True)

# set 0's to nan's as requested...
df2[df2["logRet"] == 0] = np.nan

t0 = datetime.now()

dfgp = df2.groupby(['monthEndDate', 'stock'])  # groupby object
# what you want
tot = np.exp(dfgp["logRet"].sum() - 1)
var = dfgp["logRet"].var()  # ddof=1 by default in pandas 1.1.3
vol = dfgp["volume"].mean()

print(f"{(datetime.now() - t0).total_seconds():.2f}s elapsed...")
# 2.89s elapsed...

然后,您可以根据需要使用这些数据集。例如使用pd.concat([tot, var, vol], axis=1)将它们组合在一起。

tot
Out[6]: 
monthEndDate  stock
1990-01-31    A        inf   <- will not happen with real data
              B        inf
1990-02-28    A        inf
              B        inf
Name: logRet, dtype: float64

var
Out[7]: 
monthEndDate  stock
1990-01-31    A        0.0000
              B        0.0000
1990-02-28    A        0.0025
              B        0.0025
Name: logRet, dtype: float64

vol
Out[8]: 
monthEndDate  stock
1990-01-31    A        1.5
              B        1.5
1990-02-28    A        3.5
              B        1.5
Name: volume, dtype: float64

tot部分的溢出仅仅是由于重复的增量。在真实数据中不会发生这种情况。