跨DataFrame按条件聚合高温天数并优化慢代码性能

问题背景 #

现有harvest_df（含多年收获数据，多地块对应同一气象站id）与weather_df（气象站逐日最高温数据），需为每个地块计算收获日期前2个月内最高温超30℃的天数。原使用apply逐行处理的代码耗时数分钟，需优化性能。

原低效代码：

def days_above_thresh(x, weather_df):
    return weather_df.loc[
            (weather_df["id"]==x.id) & \
            (weather_df["day"]>=x['harvest_date']-DateOffset(months=2)) & \
            (weather_df["day"]<=x['harvest_date']) & \
            (weather_df["temperature_max"]>30),
            "temperature_max"].count()

harvest_df["days_above_30"] = harvest_df.apply(days_above_thresh , args=(weather_df,), axis=1)

核心优化思路 #

抛弃逐行循环的apply，改用pandas向量化操作+分组/滚动窗口统计，利用底层C实现的运算逻辑大幅提升速度，时间复杂度从O(N*M)降至O(N+M)级别。

方案一：合并关联+分组计数（适用于中小规模数据） #

通过合并两个表的id关联数据，再筛选时间范围后分组计数，避免逐行遍历：

import pandas as pd

# 强制转换日期列为datetime类型（必须步骤）
weather_df['day'] = pd.to_datetime(weather_df['day'])
harvest_df['harvest_date'] = pd.to_datetime(harvest_df['harvest_date'])

# 计算每个收获记录的时间范围起始点（收获日前2个月）
harvest_df['start_date'] = harvest_df['harvest_date'] - pd.DateOffset(months=2)

# 给harvest_df添加临时唯一标识，用于后续分组计数
harvest_df['temp_key'] = range(len(harvest_df))

# 按id合并两个表
merged = pd.merge(harvest_df, weather_df, on='id', how='left')

# 筛选出符合时间范围且高温的记录
mask = (merged['day'] >= merged['start_date']) & \
       (merged['day'] <= merged['harvest_date']) & \
       (merged['temperature_max'] > 30)
filtered = merged[mask]

# 按临时标识分组计数，合并回原表
hot_day_counts = filtered.groupby('temp_key')['temperature_max'].count().reset_index(name='days_above_30')
harvest_df = harvest_df.merge(hot_day_counts, on='temp_key', how='left').fillna(0)

# 清理临时列
harvest_df.drop('temp_key', axis=1, inplace=True)

方案二：预计算滚动窗口+时间匹配（适用于超大规模数据） #

先为每个气象站预计算逐日过去2个月的累计高温天数，再通过高效的时间匹配关联收获记录，性能最优：

import pandas as pd

# 预处理日期列
weather_df['day'] = pd.to_datetime(weather_df['day'])
harvest_df['harvest_date'] = pd.to_datetime(harvest_df['harvest_date'])

# 标记高温天（1=高温，0=非高温）
weather_df['is_hot'] = (weather_df['temperature_max'] > 30).astype(int)

# 按id分组，补全缺失日期（避免滚动窗口计算错误）
def fill_missing_daily_data(group):
    # 生成组内完整日期序列
    full_date_range = pd.date_range(start=group['day'].min(), end=group['day'].max(), freq='D')
    # 重新索引，缺失数据补0
    return group.set_index('day')\
                .reindex(full_date_range)\
                .fillna({'is_hot': 0, 'id': group['id'].iloc[0]})\
                .reset_index()\
                .rename(columns={'index': 'day'})

weather_df_filled = weather_df.groupby('id').apply(fill_missing_daily_data).reset_index(drop=True)

# 按id分组，计算逐日过去2个月的累计高温天数
weather_df_filled['days_hot_2m'] = weather_df_filled.groupby('id')['is_hot']\
                                                    .rolling(window='60D', closed='right')\
                                                    .sum()\
                                                    .reset_index(level=0, drop=True)

# 用merge_asof高效匹配收获日期对应的高温天数
harvest_df = pd.merge_asof(
    harvest_df.sort_values('harvest_date'),
    weather_df_filled.sort_values('day'),
    left_on='harvest_date',
    right_on='day',
    by='id',
    direction='backward'  # 取不晚于收获日期的最近气象数据
)

# 整理结果列，处理缺失值
harvest_df.rename(columns={'days_hot_2m': 'days_above_30'}, inplace=True)
harvest_df['days_above_30'] = harvest_df['days_above_30'].fillna(0).astype(int)

优化效果说明 #

• 原apply逐行循环每次都要扫描全量气象数据，数据量较大时耗时呈指数增长；
• 两种优化方案均采用向量化运算，底层为C实现，速度可提升10~100倍（取决于数据规模）。

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