使用apply与lambda处理DataFrame列时出现KeyError的技术求助

问题背景 #

原有代码通过Pandas的BusinessDay偏移量，基于包含YearOfSRC和MonthNumberOfSRC列的DataFrame，生成对应年月的第7个工作日作为PredictionDate：

df['PredictionDate'] = (pd
 .to_datetime(df[['YearOfSRC', 'MonthNumberOfSRC']]
                .set_axis(['year' ,'month'], axis=1)
                .assign(day=1)
              )
 .sub(pd.offsets.BusinessDay(1))
 .add(pd.offsets.BusinessDay(7))
)

尝试改用CustomBusinessDay结合holidays包，根据美国节假日动态调整工作日时，出现KeyError: 'YearOfSRC'错误，且已确认df_final中存在该列：

df_final['PredictionDate'] = (pd
 .to_datetime(df_final[['YearOfSRC', 'MonthNumberOfSRC']]
                .set_axis(['year' ,'month'], axis=1)
                .assign(day=1)
              )
 .sub(df_final.apply(lambda x : pd.offsets.CustomBusinessDay(1, holidays = holidays.US(years = x['YearOfSRC']).keys())))
 .add(df_final.apply(lambda x: pd.offsets.CustomBusinessDay(7, holidays = holidays.US(years = x['YearOfSRC']).keys())))
)

错误原因 #

原代码中，pd.to_datetime(...)生成的是DatetimeIndex，而后续的.sub()和.add()尝试将其与df_final.apply()返回的Series做运算。Pandas会尝试对齐两者的索引，但DatetimeIndex中不存在YearOfSRC列，导致索引对齐失败抛出KeyError。同时，逐行生成CustomBusinessDay偏移量的方式也不符合Pandas的向量运算逻辑。

修正方案 #

方案一：逐行处理（简单直观） #

将整个日期计算逻辑封装到函数中，通过apply逐行处理每一条数据，确保能正确获取当前行的年份生成对应节假日规则：

import pandas as pd
import holidays

def calculate_prediction_date(row):
    # 获取当前行的年月信息
    year = row['YearOfSRC']
    month = row['MonthNumberOfSRC']
    # 生成当月第一天
    first_day = pd.to_datetime(f"{year}-{month}-01")
    # 生成对应年份的美国节假日列表
    us_holidays = holidays.US(years=year).keys()
    # 定义自定义工作日偏移量
    custom_bd = pd.offsets.CustomBusinessDay(holidays=us_holidays)
    # 计算第7个自定义工作日：上月最后一个工作日 +7个自定义工作日
    return first_day - custom_bd + custom_bd * 7

# 应用函数生成PredictionDate列
df_final['PredictionDate'] = df_final.apply(calculate_prediction_date, axis=1)

方案二：分组处理（提升效率） #

对于大数据量的DataFrame，按YearOfSRC分组处理，每个年份只生成一次节假日规则，减少重复计算：

import pandas as pd
import holidays

def process_year_group(group):
    year = group['YearOfSRC'].iloc[0]
    # 生成当前年份的美国节假日
    us_holidays = holidays.US(years=year).keys()
    custom_bd = pd.offsets.CustomBusinessDay(holidays=us_holidays)
    # 批量生成当月第一天
    first_days = pd.to_datetime({
        'year': group['YearOfSRC'],
        'month': group['MonthNumberOfSRC'],
        'day': 1
    })
    # 批量计算PredictionDate
    group['PredictionDate'] = first_days - custom_bd + custom_bd *7
    return group

# 按年份分组处理并合并结果
df_final = df_final.groupby('YearOfSRC').apply(process_year_group).reset_index(drop=True)

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