编写可对DataFrame所有日期范围列计算日期差的函数

问题场景 #

我有一个包含日期范围格式数据的DataFrame，示例如下：

import pandas as pd

df = pd.DataFrame({
    'col_1': ["'may 2021 - 2023'", "'jan 2022 - 2023'"],
    'col_2': ["'nov 2020 - feb 2021'", "'sep 2021- 2023'"]
})

目前我只能针对每一列单独写代码计算日期差，但列数不固定时需要重复修改变量，效率很低：

# 处理col_1
s = df['col_1'].str.split(r'\s*-\s*')
df['year_1'] = (pd
   .to_datetime(s.str[1])
   .sub(pd.to_datetime(s.str[0])))

# 处理col_2
t = df['col_2'].str.split(r'\s*-\s*')
df['year_2'] = (pd
   .to_datetime(t.str[1])
   .sub(pd.to_datetime(t.str[0])))

期望输出自动为每一列生成对应的日期差列，格式如下：

col_1           Year_1                     col_2            Year_2       
'may 2021 - 2023'    610 days            'nov 2020 - feb 2021'     90 days
'jan 2022 - 2023'    365 days             'sep 2021- 2023'       730 days

通用解决方案 #

可以编写一个以DataFrame为输入的函数，自动遍历所有列处理日期范围计算：

def calculate_date_diffs(df):
    processed_df = df.copy()
    for col in processed_df.columns:
        # 移除字符串两端的单引号，拆分日期范围
        split_vals = processed_df[col].str.strip("'").str.split(r'\s*-\s*', expand=True)
        # 转换为datetime类型并计算差值
        start_dates = pd.to_datetime(split_vals[0])
        end_dates = pd.to_datetime(split_vals[1])
        # 生成对应新列名
        new_col = f"Year_{col.split('_')[-1]}"
        processed_df[new_col] = end_dates - start_dates
    return processed_df

使用示例 #

# 调用函数处理数据
result_df = calculate_date_diffs(df)
print(result_df)

输出结果：

col_1                     col_2   Year_1   Year_2
0  'may 2021 - 2023'  'nov 2020 - feb 2021'  610 days  90 days
1  'jan 2022 - 2023'    'sep 2021- 2023'    365 days 730 days

函数说明 #

• 自动遍历输入DataFrame的所有列，无需手动指定列名
• 兼容带单引号的日期字符串，自动去除两端单引号
• 用正则\s*-\s*拆分日期范围，适配分隔符前后有空格的情况
• 计算日期差后生成格式为Year_列序号的新列
• 返回处理后的新DataFrame，不修改原始数据

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