拼接 Pandas 数据框的优化方式

Question

我的问题是我有一个非常大的时间序列（约 5-10 百万个 obs），其中某些事件标有标志。在这种情况下，是股票价格下跌触发了一个事件，该事件具有一个虚拟变量，如果事件被触发，则为 1 或 0。从这个时间序列中，我想提取事件本身和随后 29 天的数据。显然，这涉及到某种类型的数组拼接。

我有一些简单的代码应该可以解决问题（它只是将标志和接下来的 29 天标记为 2，但从那里过滤数据帧很简单）但它依赖于熊猫数据帧拼接，这不是很快。代码如下：

def first_drop(df):
    indexlen = len(df.dropflag[df.dropflag==1].index)

    for y in range(indexlen):                                                                                                                           
            x = df.dropflag[df.dropflag==1].index[y]
            df.dropflag[x:30]=2

    return df.dropflag

dstk['dropflag2'] = dstk[["permno","dropflag"]].groupby('permno').apply(first_drop) 

有没有其他人发现的更快的方法来进行这种类型的拼接，您可以获得下一个 x 数量的条目？我认为使用 numpy 数组或 cythonized 函数可能会更快，但我不太清楚从哪里开始。

Answer 1

这是一种可能的方法。也许没那么快，处理一个 10,000,000 行数据集需要 1 分钟。这个想法是，通过使用.shift(-i) 在随后几天用数据填充新列，它避免了循环每个groupby 内的行。它的优点是可以灵活地重塑结果数据帧，例如 stack() 以获取堆叠记录。

import pandas as pd
import numpy as np

# generate some artificial data, 10,000,000 rows
# ============================================================
np.random.seed(0)
dates = pd.date_range('2001-01-01', periods=2500, freq='B')
permno = np.arange(1000, 5000)  # 4000 symbols
multi_index = pd.MultiIndex.from_product([permno, dates], names=['permno', 'dates'])
data = np.random.randn(10000000)
dropflag = np.random.choice([0,1], size=10000000)

df = pd.DataFrame({'data': data, 'dropflag': dropflag}, index=multi_index).reset_index('permno')

Out[273]: 
            permno    data  dropflag
dates                               
2001-01-01    1000  1.7641         1
2001-01-02    1000  0.4002         1
2001-01-03    1000  0.9787         0
2001-01-04    1000  2.2409         1
2001-01-05    1000  1.8676         0
...            ...     ...       ...
2010-07-26    4999  0.5902         1
2010-07-27    4999  0.4676         1
2010-07-28    4999 -1.9447         1
2010-07-29    4999 -0.3440         1
2010-07-30    4999 -0.7402         0

[10000000 rows x 3 columns]

# processing
# ============================================================
def func(group):
    all_data = [group]
    for i in np.arange(1, 30):
        temp = group.data.shift(-i)
        temp.name = 'data_subday{}'.format(i)
        all_data.append(temp)
    dataset = pd.concat(all_data, axis=1).iloc[:-30]
    return dataset.loc[dataset.dropflag==1]

%time df.groupby('permno').apply(func)

CPU times: user 59.7 s, sys: 1.83 s, total: 1min 1s
Wall time: 1min 5s


Out[277]: 
                   permno    data  dropflag  data_subday1  data_subday2      ...        data_subday25  data_subday26  data_subday27  data_subday28  data_subday29
permno dates                                                                 ...                                                                                 
1000   2001-01-01    1000  1.7641         1        0.4002        0.9787      ...              -1.4544         0.0458        -0.1872         1.5328         1.4694
       2001-01-02    1000  0.4002         1        0.9787        2.2409      ...               0.0458        -0.1872         1.5328         1.4694         0.1549
       2001-01-04    1000  2.2409         1        1.8676       -0.9773      ...               1.5328         1.4694         0.1549         0.3782        -0.8878
       2001-01-08    1000 -0.9773         1        0.9501       -0.1514      ...               0.1549         0.3782        -0.8878        -1.9808        -0.3479
       2001-01-09    1000  0.9501         1       -0.1514       -0.1032      ...               0.3782        -0.8878        -1.9808        -0.3479         0.1563
...                   ...     ...       ...           ...           ...      ...                  ...            ...            ...            ...            ...
4999   2010-06-09    4999  2.1195         1        1.5564        1.0739      ...               0.2677         1.2637        -0.3607        -1.4011         1.1292
       2010-06-15    4999 -1.1747         1        0.2159        0.1221      ...               1.1292         1.1614         0.4842         1.3593         0.5902
       2010-06-16    4999  0.2159         1        0.1221        0.0136      ...               1.1614         0.4842         1.3593         0.5902         0.4676
       2010-06-17    4999  0.1221         1        0.0136        0.8378      ...               0.4842         1.3593         0.5902         0.4676        -1.9447
       2010-06-18    4999  0.0136         1        0.8378        0.4887      ...               1.3593         0.5902         0.4676        -1.9447        -0.3440

[4941409 rows x 32 columns]