优化的拼接熊猫数据框架方法

Question

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

我有一些简单的代码可以做到这一点（它只是将标志和接下来的29天标记为2，但是从那里可以很容易地过滤数据帧），但是它依赖于pandas数据帧拼接，这不是很快。代码如下：

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

这里有一种可能的方法。也许不是那么快，处理10000000行数据集需要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]