如何优化依赖前序值的Pandas列计算以提升性能

现有一个DataFrame，其中final_lb和final_ub列的每行计算严格依赖前一行的结果，当前使用Python for循环实现，程序性能受限于循环的解释执行开销。DataFrame结构如下：

close      basic_lb      basic_ub  final_lb  final_ub
0   18801.70  18823.800000  18823.800000       0.0       0.0
1   18794.65  18750.575000  18845.675000       0.0       0.0
2   18827.30  18758.725000  18864.925000       0.0       0.0
3   18862.10  18786.850000  18904.450000       0.0       0.0
4   18849.25  18794.340000  18911.310000       0.0       0.0
5   18850.70  18791.143500  18903.556500       0.0       0.0
6   18859.65  18798.624150  18907.625850       0.0       0.0
7   18840.80  18795.519235  18904.930765       0.0       0.0
8   18831.95  18782.764812  18887.835188       0.0       0.0
9   18847.00  18787.678330  18891.271670       0.0       0.0
10  18815.00  18772.292997  18888.957003       0.0       0.0
11  18812.30  18759.911198  18870.938802       0.0       0.0
12  18816.95  18752.947578  18863.402422       0.0       0.0
13  18817.60  18762.375320  18869.674680       0.0       0.0
14  18815.15  18764.835288  18864.914712       0.0       0.0

当前使用的循环代码：

for i in range(1, len(data)):
    data['final_ub'].iat[i] = data['basic_ub'].iat[i] if data['basic_ub'].iat[i] < data['final_ub'].iat[i - 1] or data['close'].iat[i - 1] > data['final_ub'].iat[i - 1] else data['final_ub'].iat[i - 1]
    data['final_lb'].iat[i] = data['basic_lb'].iat[i] if data['basic_lb'].iat[i] > data['final_lb'].iat[i - 1] or data['close'].iat[i - 1] < data['final_lb'].iat[i - 1] else data['final_lb'].iat[i - 1]

1. 使用Numba JIT编译加速循环 #

Python原生for循环的性能瓶颈在于解释执行的开销，Numba可将Python函数编译为机器码，大幅提升循环效率，且逻辑与原代码完全一致，是最推荐的方案。

代码示例： #

import numba as nb
import numpy as np

@nb.njit  # 编译为机器码，默认禁用并行（因计算依赖前序结果）
def calculate_bounds(close, basic_lb, basic_ub, initial_lb, initial_ub):
    n = len(close)
    final_lb = np.zeros(n, dtype=np.float64)
    final_ub = np.zeros(n, dtype=np.float64)
    final_lb[0] = initial_lb
    final_ub[0] = initial_ub
    
    for i in range(1, n):
        # 计算final_ub[i]
        if basic_ub[i] < final_ub[i-1] or close[i-1] > final_ub[i-1]:
            final_ub[i] = basic_ub[i]
        else:
            final_ub[i] = final_ub[i-1]
        
        # 计算final_lb[i]
        if basic_lb[i] > final_lb[i-1] or close[i-1] < final_lb[i-1]:
            final_lb[i] = basic_lb[i]
        else:
            final_lb[i] = final_lb[i-1]
    
    return final_lb, final_ub

# 传入numpy数组计算（Numba对numpy支持更优）
final_lb, final_ub = calculate_bounds(
    data['close'].values,
    data['basic_lb'].values,
    data['basic_ub'].values,
    data['final_lb'].iloc[0],
    data['final_ub'].iloc[0]
)

# 赋值回DataFrame
data['final_lb'] = final_lb
data['final_ub'] = final_ub

优势： #

• 编译后性能接近原生C/C++水平，百万级数据量下速度可提升几十至上百倍
• 无需修改原计算逻辑，迁移成本低

2. Cython手动编译循环（极致性能需求） #

若Numba无法满足极致性能要求，可使用Cython编写循环代码，手动指定变量类型进一步优化。但需学习Cython语法，复杂度较高，仅适合核心性能瓶颈场景。

3. 结果验证 #

运行优化代码后，可通过以下方式验证结果与原循环一致：

# 保存原循环结果
data_old = data.copy()
for i in range(1, len(data_old)):
    data_old['final_ub'].iat[i] = data_old['basic_ub'].iat[i] if data_old['basic_ub'].iat[i] < data_old['final_ub'].iat[i - 1] or data_old['close'].iat[i - 1] > data_old['final_ub'].iat[i - 1] else data_old['final_ub'].iat[i - 1]
    data_old['final_lb'].iat[i] = data_old['basic_lb'].iat[i] if data_old['basic_lb'].iat[i] > data_old['final_lb'].iat[i - 1] or data_old['close'].iat[i - 1] < data_old['final_lb'].iat[i - 1] else data_old['final_lb'].iat[i - 1]

# 对比结果
print(data[['final_lb', 'final_ub']].equals(data_old[['final_lb', 'final_ub']]))  # 应输出True

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