Python Pandas 多索引数据帧的高效合并

Question

我有两个数据框：第一个是包含两只股票收益的多指数框架，用以下方式表示：

import pandas as pd
from pandas import IndexSlice as idx
import numpy as np

dates = pd.date_range('2000-12-31', periods=6, freq='M', name='Date')
arrays = [dates.tolist()+dates.tolist(),["10000"]*6+["10001"]*6]
index = pd.MultiIndex.from_tuples(list(zip(*arrays)), names=["Date", "Stock"])
df1 = pd.Series(np.random.randn(12), index=index).to_frame('Return').sort_index()

第二帧表示为：

并包含给定时间跨度的股票代码。

data = {'Stock':['10000','10000','10000','10001'],
        'Start':['1990-12-31', '2001-03-05', '2001-05-19', '1991-03-31'],
        'End':['2001-03-04', '2001-05-18', '2002-01-31', '2001-04-03'],
        'Code':['10','11','10','10']}
df2 = pd.DataFrame(data)
df2 = df2.set_index('Stock').sort_index()
df2['Start'] = pd.to_datetime(df2['Start']) 
df2['End'] = pd.to_datetime(df2['End']) 

各个跨度的开始日期在“开始”列中给出，结束日期在“结束”列中给出。我想将股票代码（在 df2 的“代码”列中给出）添加到 df1，这样如果股票的时间索引（df1）在 df2 的“开始”和“结束”列之间，则相应的条目在 df1 的“代码”列中包含 df2 中给出的代码。如果没有指定代码或df2中没有时间跨度覆盖df1中的时间索引，则应分配np.nan。

考虑示例 Date='2001-03-31' 和 Stock='10000'。根据 df2 的第二行，我们看到 Stock 10000 从 2001-03-05 到 2001-05-18 的代码为 11。因此，行中的条目 ('2001-03-31','10000')并且在 df1 的“代码”列中应包含 11。

此处显示了整个 df1：

以下函数具有所需的功能，但（即使是并行化）非常慢：

df1s = df1.swaplevel().sort_index().copy() # reorder s.t. date is first
df1['Code'] = np.nan
for p_tmp in df2.index.drop_duplicates().values:
    d_tmp = df1s.loc[idx[p_tmp,:]].index.get_level_values(0) # End of each month
    output = np.array([np.nan]*df1s.loc[idx[p_tmp,:]].index.size)
    if isinstance(df2.loc[p_tmp], pd.Series):
        b_mask = np.array((d_tmp >= pd.to_datetime(df2.loc[p_tmp]['Start'])) & (d_tmp <= pd.to_datetime(df2.loc[p_tmp]['End'])))
        output[b_mask] = df2.loc[p_tmp]['Code']
    else:
        for index, row in df2.loc[p_tmp].iterrows():
            b_mask = np.array((d_tmp >= pd.to_datetime(row['Start'])) & (d_tmp <= pd.to_datetime(row['End'])))
            output[b_mask] = row['Code']
    df1s.loc[p_tmp,'Code'] = output.copy()
df1 = df1s.swaplevel().sort_index().copy()

请有人给我一些提示，我可以如何实现加速？ :)

Answer 1

解决方案可能是这样的：

df_tmp = pd.DataFrame(index=df1.index).join(df2, on='Stock', how='left')
blist_validlink = (df_tmp['Start']<=df_tmp.index.get_level_values('Date'))&(df_tmp['End']>=df_tmp.index.get_level_values('Date')) 
df_tmp = df_tmp[blist_validlink]
df1 = df1.join(df_tmp['Code'], on=['Date','Stock'])