如何向量化Pandas双层循环以提升大DataFrame运行效率？

问题描述 #

我写了一段Pandas双层for循环代码，功能是遍历目标DataFrame（df）的所有列与时间戳索引，匹配另一数据集（dat）中满足start < 时间戳、end > 时间戳且plant_type与df列名一致的记录，计算这些记录的Difference总和后赋值给df对应位置。代码能正常运行，但因df数据量较大，单次运行耗时长达11分钟。我尝试用向量化方式优化，但多次尝试均报错无法执行，恳请提供优化方案。

数据结构 #

df结构 #

{'Fossil Gas': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Fossil Brown coal/Lignite': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Biomass': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Nuclear': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Fossil Oil': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Wind Onshore': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Solar': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Fossil Hard coal': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}, 'Other renewable': {Timestamp('2022-01-01 00:00:00', freq='H'): 0, Timestamp('2022-01-01 01:00:00', freq='H'): 0, Timestamp('2022-01-01 02:00:00', freq='H'): 0, Timestamp('2022-01-01 03:00:00', freq='H'): 0, Timestamp('2022-01-01 04:00:00', freq='H'): 0}}

dat结构 #

{'Unnamed: 0': {Timestamp('2016-07-19 12:00:00'): 0, Timestamp('2018-01-01 00:00:00'): 243}, 'avail_qty': {Timestamp('2016-07-19 12:00:00'): 0.0, Timestamp('2018-01-01 00:00:00'): 0.0}, 'biddingzone_domain': {Timestamp('2016-07-19 12:00:00'): 'HU', Timestamp('2018-01-01 00:00:00'): 'HU'}, 'businesstype': {Timestamp('2016-07-19 12:00:00'): 'Unplanned outage', Timestamp('2018-01-01 00:00:00'): 'Planned maintenance'}, 'curvetype': {Timestamp('2016-07-19 12:00:00'): 'A03', Timestamp('2018-01-01 00:00:00'): 'A03'}, 'docstatus': {Timestamp('2016-07-19 12:00:00'): nan, Timestamp('2018-01-01 00:00:00'): nan}, 'end': {Timestamp('2016-07-19 12:00:00'): Timestamp('2018-01-25 00:00:00'), Timestamp('2018-01-01 00:00:00'): Timestamp('2019-01-01 00:00:00')}, 'mrid': {Timestamp('2016-07-19 12:00:00'): 'doEAXAnlBP-GXBLvSXLrww', Timestamp('2018-01-01 00:00:00'): 'm7N8fbZN3jM4LLrCTgu1zw'}, 'nominal_power': {Timestamp('2016-07-19 12:00:00'): 25.0, Timestamp('2018-01-01 00:00:00'): 95.0}, 'plant_type': {Timestamp('2016-07-19 12:00:00'): 'Fossil Gas', Timestamp('2018-01-01 00:00:00'): 'Fossil Gas'}, 'production_resource_id': {Timestamp('2016-07-19 12:00:00'): '15WDUME------PPI', Timestamp('2018-01-01 00:00:00'): '15WDKCE------PPQ'}, 'production_resource_location': {Timestamp('2016-07-19 12:00:00'): 'Százhalombatta', Timestamp('2018-01-01 00:00:00'): 'Debrecen'}, 'production_resource_name': {Timestamp('2016-07-19 12:00:00'): 'Dunamenti Erőmű', Timestamp('2018-01-01 00:00:00'): 'Debreceni Kombináltciklusú Erőmű'}, 'pstn': {Timestamp('2016-07-19 12:00:00'): 1, Timestamp('2018-01-01 00:00:00'): 1}, 'qty_uom': {Timestamp('2016-07-19 12:00:00'): 'MAW', Timestamp('2018-01-01 00:00:00'): 'MAW'}, 'resolution': {Timestamp('2016-07-19 12:00:00'): 'PT60M', Timestamp('2018-01-01 00:00:00'): 'PT15M'}, 'revision': {Timestamp('2016-07-19 12:00:00'): 1, Timestamp('2018-01-01 00:00:00'): 944}, 'start': {Timestamp('2016-07-19 12:00:00'): Timestamp('2016-07-19 12:00:00'), Timestamp('2018-01-01 00:00:00'): Timestamp('2018-01-01 00:00:00')}, 'created_doc_time': {Timestamp('2016-07-19 12:00:00'): '2018-01-25 06:00:54+01:00', Timestamp('2018-01-01 00:00:00'): '2018-12-30 15:36:00+01:00'}, 'Deleted': {Timestamp('2016-07-19 12:00:00'): 25.0, Timestamp('2018-01-01 00:00:00'): 95.0}}

优化方案建议 #

注：以下方案中假设Difference对应dat中的Deleted列，若实际为其他列可自行替换。

方法1：长格式合并+分组求和 #

将df转为长格式后与dat合并，筛选时间匹配的记录再分组求和，最后转回宽格式更新原df，全程避免循环：

import pandas as pd

# 将df转为长格式（时间戳+plant_type为行维度）
df_long = df.stack().reset_index()
df_long.columns = ['timestamp', 'plant_type', 'value']

# 重置dat索引（原索引无业务意义）
dat_reset = dat.reset_index(drop=True)

# 按plant_type合并两个数据集
merged = pd.merge(df_long, dat_reset, on='plant_type', how='left')

# 筛选时间戳在[start, end)区间内的记录
mask = (merged['timestamp'] > merged['start']) & (merged['timestamp'] < merged['end'])
filtered = merged[mask]

# 按时间戳和plant_type分组求和
summed = filtered.groupby(['timestamp', 'plant_type'])['Deleted'].sum().reset_index()

# 转回宽格式并更新原df
result = summed.pivot(index='timestamp', columns='plant_type', values='Deleted').fillna(0)
df.update(result)

方法2：时间区间匹配 #

利用Pandas的IntervalIndex快速判断时间戳是否在区间内，减少循环次数：

# 为dat创建时间区间列（左开右闭，匹配需求中的start<时间戳<end）
dat['time_interval'] = pd.IntervalIndex.from_arrays(dat['start'], dat['end'], closed='neither')

# 仅按plant_type循环，内部用矢量化判断
for plant in df.columns:
    plant_records = dat[dat['plant_type'] == plant]
    if plant_records.empty:
        continue
    # 对每个时间戳，计算符合条件的Deleted总和
    df[plant] = df.index.to_series().apply(
        lambda ts: plant_records[plant_records['time_interval'].contains(ts)]['Deleted'].sum()
    )

方法3：NumPy广播矢量化 #

将时间戳和区间转为NumPy数组，用广播机制批量计算匹配情况，性能最优：

import numpy as np

# 提取df的时间戳数组和plant类型列表
timestamps = df.index.values
plants = df.columns.values

for plant in plants:
    plant_data = dat[dat['plant_type'] == plant]
    if plant_data.empty:
        continue
    # 提取dat中的关键数组
    starts = plant_data['start'].values
    ends = plant_data['end'].values
    deletes = plant_data['Deleted'].values
    
    # 广播判断每个时间戳是否在每个区间内
    match_mask = (timestamps[:, np.newaxis] > starts) & (timestamps[:, np.newaxis] < ends)
    # 按时间轴求和
    total_deletes = np.sum(match_mask * deletes, axis=1)
    # 赋值给df对应列
    df[plant] = total_deletes

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