如何在Pandas中利用date_range扩展多层索引？

输入DataFrame #

我们有如下输入DataFrame，使用Pandas 1.3.2生成的代码如下：

import pandas as pd
import datetime

input_data = [
    ["1", datetime.datetime(2023,2,21,20,0,0), 10],
    ["1", datetime.datetime(2023,2,21,20,30,0), 10],
    ["2", datetime.datetime(2023,2,21,15,0,0), 15],
    ["2", datetime.datetime(2023,2,21,15,30,0), 15],
]

df_input = pd.DataFrame(data=input_data, columns=["id", "time", "duration"]).set_index(["id", "time"])

需求说明 #

需要根据duration列的时长，扩展DataFrame的第二层索引（time列）。以id为"1"的数据为例：要生成从20:00到20:30的所有10分钟间隔时间槽（即20:00、20:10、20:20、20:30），最终输出的DataFrame需包含这些扩展后的时间索引。

现有低效实现 #

目前已实现的解决方案运行速度较慢，代码如下：

import pandas as pd
import datetime

input_data = [
    ["1", datetime.datetime(2023,2,21,20,0,0), 10],
    ["1", datetime.datetime(2023,2,21,20,30,0), 10],
    ["2", datetime.datetime(2023,2,21,15,0,0), 15],
    ["2", datetime.datetime(2023,2,21,15,30,0), 15],
]

df_input = pd.DataFrame(data=input_data, columns=["id", "time", "duration"]).set_index(["id", "time"])

df_output = pd.DataFrame()

for i in range(0, df_input.shape[0], 2):
    start_at = df_input.index[i][1]
    end_at = df_input.index[i+1][1]
    duration = df_input.iloc[i]["duration"]
    
    df_cut = pd.DataFrame(
        pd.date_range(
            start=start_at,
            end=end_at,
            freq=f"{duration}min",
        ).rename("start_at_converted")
    )
    df_cut["id"] = df_input.index[i][0]
    df_cut["duration"] = duration
    
    df_output = pd.concat((df_output, df_cut), axis=0)

df_output = df_output.set_index(["id", "start_at_converted"])

高效解决方案 #

可以利用Pandas的分组操作和向量化方法替代循环，大幅提升效率（尤其在数据量较大时）。具体实现代码如下：

import pandas as pd
import datetime

input_data = [
    ["1", datetime.datetime(2023,2,21,20,0,0), 10],
    ["1", datetime.datetime(2023,2,21,20,30,0), 10],
    ["2", datetime.datetime(2023,2,21,15,0,0), 15],
    ["2", datetime.datetime(2023,2,21,15,30,0), 15],
]

df_input = pd.DataFrame(data=input_data, columns=["id", "time", "duration"]).set_index(["id", "time"])

# 重置索引，将id和time转为普通列以便分组
df_reset = df_input.reset_index()

# 按id分组生成时间序列并展开
df_output = df_reset.groupby("id").apply(
    lambda x: pd.DataFrame({
        "start_at_converted": pd.date_range(
            start=x["time"].iloc[0],
            end=x["time"].iloc[-1],
            freq=f"{x['duration'].iloc[0]}min"
        ),
        "duration": x["duration"].iloc[0]
    })
).reset_index().set_index(["id", "start_at_converted"]).drop(columns="level_1")

print(df_output)

优势说明 #

• 避免了循环中反复调用pd.concat的性能损耗（循环内concat会不断创建新对象，效率极低）；
• 利用Pandas分组操作groupby结合apply，内部采用向量化处理，比Python原生循环效率提升显著；
• 代码逻辑更清晰，结构更简洁，易于维护。

内容的提问来源于stack exchange，提问作者Florian Drouet