如何优化Pandas DataFrame批量比值计算的运行性能？

问题说明 #

现有两个Pandas DataFrame：

import pandas as pd
import numpy as np

bdata = pd.DataFrame(
    {
        'A': [np.nan, np.nan, np.nan, np.nan, np.nan, 92.0, 49.0, 56.0, 72.0, np.nan],
        'B': [np.nan, np.nan, np.nan, np.nan, np.nan, 69.0, 60.0, np.nan, np.nan, np.nan],
        'C': [np.nan, np.nan, np.nan, np.nan, 97.0, 20.0, 56.0, 95.0, np.nan, 74.0],
        'D': [np.nan, np.nan, np.nan, np.nan, np.nan, 75.0, np.nan, 15.0, 84.0, np.nan]
    },
    index=pd.date_range('2018-01-01', periods=10)
)

sdata = pd.DataFrame(
    {
        'A': [np.nan, np.nan, np.nan, np.nan, 7.0, np.nan, np.nan, np.nan, np.nan, np.nan],
        'B': [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 5.0],
        'C': [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
        'D': [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 3.0, np.nan, np.nan, 0.0]
    },
    index=pd.date_range('2018-01-01', periods=10)
)

需要计算的规则：
对每个日期idate，遍历每一列：

1. 取bdata中从idate开始的第一个非NaN值（记为b_val）
2. 取sdata中索引大于b_val所在索引的第一个非NaN值（记为s_val）
3. 计算比值s_val / b_val，最终将所有日期的所有列比值求和

示例：

• 2018-01-05的比值：

  A          B          C          D
  NaN/92.0   5.0/69.0   NaN/97.0   3.0/75.0
  

• 2018-01-06的比值：

  A          B          C          D
  NaN/92.0   5.0/69.0   NaN/20.0   3.0/75.0
  

原代码采用逐日期循环处理，在数据量较大（数千日期+多列）时速度极慢：

ratiosum=0
for idate in bdata.index[:-1]:
    bdata1=bdata.loc[bdata.index>=idate].fillna(method='ffill')
    sdata1=sdata.loc[bdata.index>=idate][bdata1.shift(1)>0] #bdata is always >0
    bvalue1=bdata1.fillna(method='bfill').iloc[0]
    svalue1=sdata1.fillna(method='bfill').iloc[0]
    ratiosum +=(svalue1/bvalue1).sum() #Actually, I only want the sum of all ratios

优化方案 #

核心思路 #

避免逐日期循环，利用Pandas矢量化操作+二分查找，批量预处理每列的匹配关系，一次性计算所有比值的总和。

优化代码 #

import pandas as pd
import numpy as np
from bisect import bisect_right

def calculate_total_ratio(bdata, sdata):
    total = 0.0
    date_index = bdata.index
    n_dates = len(date_index)
    
    for col in bdata.columns:
        # 提取当前列bdata的非NaN数据：(索引位置, 值)
        b_non_nan = bdata[col].dropna()
        if b_non_nan.empty:
            continue
        b_positions = date_index.get_indexer(b_non_nan.index)
        b_values = b_non_nan.values
        
        # 提取当前列sdata的非NaN数据：(索引位置, 值)
        s_non_nan = sdata[col].dropna()
        if s_non_nan.empty:
            # 没有匹配的s值，所有比值为NaN，不参与求和
            continue
        s_positions = date_index.get_indexer(s_non_nan.index)
        s_values = s_non_nan.values
        
        # 为每个bdata的非NaN值，找到对应的sdata中索引更大的第一个非NaN值
        s_matches = []
        for b_pos in b_positions:
            # 找到第一个大于b_pos的s_pos的索引
            idx = bisect_right(s_positions, b_pos)
            if idx < len(s_values):
                s_matches.append(s_values[idx])
            else:
                s_matches.append(np.nan)
        
        # 计算每个bdata值对应的日期区间，以及该区间内的比值总和
        prev_pos = -1
        for i in range(len(b_positions)):
            current_pos = b_positions[i]
            # 确定当前区间的日期范围（对应原代码的index[:-1]）
            start = prev_pos + 1
            end = min(current_pos, n_dates - 2)
            if start > end:
                break
            count = end - start + 1
            
            # 累加该区间的比值总和
            if not np.isnan(s_matches[i]):
                ratio = s_matches[i] / b_values[i]
                total += ratio * count
            
            prev_pos = current_pos
        
        # 处理最后一个bdata值之后的剩余日期
        start = prev_pos + 1
        if start <= n_dates - 2:
            count = (n_dates - 2) - start + 1
            if not np.isnan(s_matches[-1]):
                ratio = s_matches[-1] / b_values[-1]
                total += ratio * count
    
    return total

# 调用计算
ratiosum = calculate_total_ratio(bdata, sdata)
print(ratiosum)

性能提升说明 #

1. 减少循环次数：原代码循环每个日期，优化后仅循环每个列，循环次数从数千次降至列数（通常远小于日期数）。
2. 二分查找加速匹配：用bisect_right快速定位sdata中符合条件的第一个值，时间复杂度为O(log S)（S为列内sdata非NaN数量），远快于逐行查找。
3. 批量区间求和：直接计算每个bdata值覆盖的日期区间长度，一次性累加该区间的比值总和，避免重复计算每个日期的比值。

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