如何获取Pandas DataFrame对应年月的最大天数并优化运行效率？

原始数据与需求 #

我们有如下Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({'ID': [1,2,3,4],
                   'month': [12, 2, 8, 6],
                   'year': [2021, 2022, 2022, 2020]})

输出结果：

ID  month  year
0   1     12  2021
1   2      2  2022
2   3      8  2022
3   4      6  2020

需要为该DataFrame新增Max_day列，存储每行对应年月的最大天数。

原始低效实现 #

最初使用循环逐行计算的代码如下，虽然能得到正确结果，但数据量大时运行耗时极长：

df['Max_day'] = 0

for i in range(0,len(df)):
    v_month = df['month'].loc[i]
    v_year = df['year'].loc[i]

    df['Max_day'].loc[i] = pd.Period(v_year.astype(str) + "-" + 
                                     v_month.astype(str)).days_in_month

正确输出结果：

ID  month  year  Max_day
0   1     12  2021       31
1   2      2  2022       28
2   3      8  2022       31
3   4      6  2020       30

高效优化方案 #

Pandas的核心优势是向量化操作，完全可以避免逐行循环，以下两种方法能大幅提升效率：

方法1：利用pd.to_datetime提取月天数

通过构造年月字符串转为datetime类型，直接调用dt.days_in_month属性获取每月最大天数：

df['Max_day'] = pd.to_datetime(
    df['year'].astype(str) + '-' + df['month'].astype(str),
    format='%Y-%m'
).dt.days_in_month

方法2：直接构造PeriodIndex

使用PeriodIndex直接批量创建年月周期对象，再提取days_in_month：

df['Max_day'] = pd.PeriodIndex(
    year=df['year'], 
    month=df['month'], 
    freq='M'
).days_in_month

优化原理 #

循环逐行处理时，每次loc索引查找和单个Period对象创建都会带来额外开销；而向量化操作是底层基于C语言的批量处理，能一次性完成所有行的计算，效率提升非常显著（数据量越大，差距越明显）。

内容的提问来源于stack exchange，提问作者valentim.kodak