Pandas shift datetimeindex 运行时间过长

Question

我在使用日期时间索引移动大型数据框时遇到了运行时间问题。

使用创建的虚拟数据的示例：

df = pd.DataFrame({'col1':[0,1,2,3,4,5,6,7,8,9,10,11,12,13]*10**5,'col3':list(np.random.randint(0,100000,14*10**5)),'col2':list(pd.date_range('2020-01-01','2020-08-01',freq='M'))*2*10**5})
df.col3=df.col3.astype(str)
df.drop_duplicates(subset=['col3','col2'],keep='first',inplace=True)

如果我不使用 datetimeindex 换档，只需要大约 12 秒：

%%time
tmp=df.groupby('col3')['col1'].shift(2,fill_value=0)
Wall time: 12.5 s

但是当我使用 datetimeindex 时，作为我需要的那种情况，大约需要 40 分钟：

%%time    
tmp=df.set_index('col2').groupby('col3')['col1'].shift(2,freq='M',fill_value=0)
Wall time: 40min 25s

在我的情况下，我需要从 shift(1) 到 shift(6) 的数据，并将它们与 col2 和 col3 的原始数据合并。所以我使用for 循环和合并。 有什么解决办法吗？感谢您的回答，非常感谢您的回复。

Ben 的回答解决了这个问题：

%%time
tmp=df1[['col1','col3', 'col2']].assign(col2 = lambda x: x['col2'] + MonthEnd(2)).set_index(['col3', 'col2']).add_suffix(f'_{2}').fillna(0).reindex(pd.MultiIndex.from_frame(df1[['col3','col2']])).reset_index()
Wall time: 5.94 s

也实现了循环：

%%time
res=(pd.concat([df1.assign(col2 = lambda x: x['col2'] + MonthEnd(i)).set_index(['col3', 'col2']).add_suffix(f'_{i}') for i in range(0,7)],axis=1).fillna(0)).reindex(pd.MultiIndex.from_frame(df1[['col3','col2']])).reset_index() 
Wall time: 1min 44s

实际上，我的真实数据已经在使用MonthEnd(0)，所以我只是在range(1,7) 中使用循环。我还实现了多列，所以我不使用astype，并实现reindex，因为我使用left merge。

Answer 1

这两个操作略有不同，结果也不相同，因为您的数据（至少这里的虚拟数据）没有排序，特别是如果您缺少某些 col3 值的日期。也就是说，时差似乎很大。所以我认为你应该采取一些不同的方式。

一种方法是将 X MonthEnd 添加到 col2 以使 X 从 0 到 6，使用 concat 所有这些，在 set_index 之后 col3 和 col2，add_suffix 以跟踪“移位”值. fillna 并将dtype 转换为原始值。其余的主要是化妆品，具体取决于您的需求。

from pandas.tseries.offsets import MonthEnd

res = (
    pd.concat([
        df.assign(col2 = lambda x: x['col2']  + MonthEnd(i))
          .set_index(['col3', 'col2'])
          .add_suffix(f'_{i}')
        for i in range(0,7)], 
        axis=1)
      .fillna(0) 
      # depends on your original data
      .astype(df['col1'].dtype) 
      # if you want a left merge ordered like original df
      #.reindex(pd.MultiIndex.from_frame(df[['col3','col2']]))
      # if you want col2 and col3 back as columns
      # .reset_index() 
)

请注意，concat 默认情况下会进行外部连接，因此您最终会得到月份，该月份不在您的原始数据中，而 col1_0 实际上是带有我的随机数的原始数据。

print(res.head(10))
                 col1_0  col1_1  col1_2  col1_3  col1_4  col1_5  col1_6
col3 col2                                                              
0    2020-01-31       7       0       0       0       0       0       0
     2020-02-29       8       7       0       0       0       0       0
     2020-03-31       2       8       7       0       0       0       0
     2020-04-30       3       2       8       7       0       0       0
     2020-05-31       4       3       2       8       7       0       0
     2020-06-30      12       4       3       2       8       7       0
     2020-07-31      13      12       4       3       2       8       7
     2020-08-31       0      13      12       4       3       2       8
     2020-09-30       0       0      13      12       4       3       2
     2020-10-31       0       0       0      13      12       4       3

Answer 2

这是groupby + shift 的问题。问题是，如果您指定的轴不是0 或频率它falls back to a very slow loop over the groups。如果两者都没有指定，它可以使用更快的路径，这就是为什么您会看到性能之间存在数量级差异的原因。

DataFrame.GroupBy.shift中的相关代码为：

def shift(self, periods=1, freq=None, axis=0, fill_value=None):
    """..."""
    if freq is not None or axis != 0:
        return self.apply(lambda x: x.shift(periods, freq, axis, fill_value))

---

之前这个问题扩展到指定fill_value