基于Pandas计算合同各周期生效天数的优化方案问询

需求说明 #

我有一个存储合同起止日期的Pandas DataFrame，需要计算合同所覆盖的所有周期（例如月度）内的合同生效天数。

示例输入 #

start_date   end_date
0 2022-01-01 2022-02-15
1 2022-02-01 2022-04-01
2 2022-03-01 2022-04-15

预期输出 #

2022-01    30
2022-02    41
2022-03    61
2022-04    14
Freq: M, dtype: int64

现状与诉求 #

我已实现一个可行的朴素方案，但因需处理数百万行数据，效率至关重要，希望获得更高效、更贴合Pandas/Python风格的优化建议。

当前方案确定最小覆盖周期集合后，会脱离数组函数，改用行与周期的循环遍历。我曾寻找可计算周期重叠或时间差的数组函数，但仅找到start_time和end_time可用。

原始实现代码 #

import pandas


def days_in_periods(df: pandas.DataFrame, inc_st: bool = True, inc_en: bool = True, period_freq='M') -> pandas.Series:
    """ Calculate the days in each period covered by any contract defined within the dataframe """
    # 创建周期范围
    periods = pandas.period_range(start=df['start_date'].min(),
                                  end=df['end_date'].max(),
                                  freq=period_freq)

    period_days = pandas.Series(data=[0] * len(periods),
                                index=periods,
                                dtype=int)

    for index, row in df.iterrows():
        st = row['start_date']
        en = row['end_date']

        print(f'contract: {st:%d/%m} - {en:%d/%m}')

        total_days: int = (en - st).days + inc_en - (1 - inc_st)

        print(f'contract days: {total_days}')

        total_days_check: int = 0

        for period in periods:
            per_st = period.start_time
            per_en = period.end_time
            print(f'\tperiod: {per_st:%d/%m} - {per_en:%d/%m}', end='')

            if per_en < st or per_st > en:
                print('\t0')
                continue

            days: int = (per_en - per_st).days + 1

            if per_st <= st <= per_en:
                days -= (st - per_st).days + (1 - inc_st)

            if per_st <= en <= per_en:
                days -= (per_en - en).days + (1 - inc_en)

            total_days_check += days

            print(f'\t{days}')

            period_days[period] += days

        print(f'total days check: {total_days_check}')
        assert total_days == total_days_check

    return period_days


# 创建示例DataFrame
df_ex = pandas.DataFrame({'start_date': ['2022-01-01', '2022-02-01', '2022-03-01'],
                          'end_date': ['2022-02-15', '2022-04-01', '2022-04-15']})

# 转换日期列为datetime类型
df_ex['start_date'] = pandas.to_datetime(df_ex['start_date'])
df_ex['end_date'] = pandas.to_datetime(df_ex['end_date'])

days_in_periods(df_ex, inc_st=True, inc_en=True)
days_in_periods(df_ex, inc_st=True, inc_en=False)
days_in_periods(df_ex, inc_st=False, inc_en=True)
print(days_in_periods(df_ex, inc_st=False, inc_en=False))

经优化后的代码（基于sammywemmy建议） #

import operator
import pandas


def days_in_periods(df: pandas.DataFrame,
                    inc_st: bool = True,
                    inc_en: bool = True,
                    period_freq='M') -> pandas.DataFrame:
    """ Calculate the days in each period covered by any contract defined within the dataframe """
    # 生成所有覆盖日期的序列
    day_range = pandas.date_range(df['start_date'].min(),
                                  df['end_date'].max(),
                                  freq='D').to_series(name='days').reset_index(drop=True)

    # 根据是否包含起止日期选择比较运算符
    st_op = operator.le if inc_st else operator.lt
    en_op = operator.ge if inc_en else operator.gt

    # 交叉连接合同与日期，筛选生效日期
    df = df.merge(day_range, how='cross')
    df = (df.loc[st_op(df['start_date'], df['days']) & en_op(df['end_date'], df['days'])]
            .resample(on='days', rule=period_freq)
            .size()
          )
    # 将索引转换为周期类型
    df.index = df.index.to_period(period_freq)

    return df


# 创建示例DataFrame
df_ex = pandas.DataFrame({'start_date': ['2022-01-01', '2022-02-01', '2022-03-01'],
                          'end_date': ['2022-02-15', '2022-04-01', '2022-04-15']})

# 转换日期列为datetime类型
df_ex['start_date'] = pandas.to_datetime(df_ex['start_date'])
df_ex['end_date'] = pandas.to_datetime(df_ex['end_date'])

print(days_in_periods(df_ex, inc_st=False, inc_en=False))

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