Pandas shift datetimeindex 花费的时间太长运行

Pandas shift datetimeindex takes too long time running

我在移动具有日期时间索引的大型数据帧时遇到 运行 时间问题。

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

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)

如果我shift不使用datetimeindex,只需要12s左右:

%%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) 的数据,并通过 col2col3 将它们与原始数据合并。所以我使用 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.

这两个操作略有不同,结果也不相同,因为您的数据(至少这里的虚拟数据)没有排序,尤其是当您缺少某些 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

这是 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