Pandas DataFrame基于前一行滞后值的逐行函数优化需求

我需要在Pandas DataFrame中逐行调用自定义函数，该函数依赖前一行计算得到的Q和S的滞后值，首行已具备Q和S的初始值，从第二行开始计算。当前用for循环能正常运行，但目标DataFrame有3000多行，循环效率不够。试过df.shift(-1)、rolling.apply()和向量化方法，都没成功实现。

原始代码如下：

import time
import pandas as pd
import math

def myfunc(Eo, P, Smax, Sprev, Qprev):
  print("i = ", i)
  print("Qprev = ", Qprev)
  S = Sprev + Eo * math.exp(-1 * Sprev/Smax) - P + Qprev
  Q = P + S
  print("Q = ", Q)
  return S, Q

data = {'peti': {0: 0.1960418075323104, 1: 0.5796640515327454, 2: 0.737823486328125, 3: 0.222676545381546, 4: 0.8804306983947754}, 'tas': {0: 281.0088195800781, 1: 277.112060546875, 2: 273.7044372558594, 3: 277.48309326171875, 4: 279.4878845214844}, 'precip': {0: 0.0, 1: 0.0, 2: 1.5046296539367177e-05, 3: 0.0002500000118743, 4: 4.6296295295178425e-06}, 'year': {0: 2008, 1: 2008, 2: 2008, 3: 2008, 4: 2008}, 'row_id': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4}, 'S': {0: 90.9, 1: "nan", 2: "nan", 3: "nan", 4: "nan"}, 'Q': {0: 0.0, 1: "nan", 2: "nan", 3: "nan", 4: "nan"}}
    
df = pd.DataFrame.from_dict(data)

smaxval = 100

start_time = time.time()

for i in df.index[1:len(df)]: # 从第二行开始
  df.loc[i,["S","Q"]] = myfunc(
    df.peti[i],
      df.precip[i],
        smaxval,
          df.S[i-1],
            df.Q[i-1])
            
print("--- %s seconds ---" % (time.time() - start_time))

你的计算属于强递推依赖（每一步的S和Q完全依赖上一步的计算结果），这种场景下普通向量化方法无法实现，最有效的优化方式是用JIT编译加速循环，numba库可以将Python循环编译为机器码，大幅提升运行效率。

优化代码实现 #

import time
import pandas as pd
import numpy as np
from numba import jit

# 用numba编译计算函数，nopython=True表示完全编译为机器码
@jit(nopython=True)
def compute_s_q(peti_arr, precip_arr, smax_val, init_S, init_Q):
    n = len(peti_arr)
    S_arr = np.zeros(n)
    Q_arr = np.zeros(n)
    # 初始化首行值
    S_arr[0] = init_S
    Q_arr[0] = init_Q
    
    for i in range(1, n):
        Sprev = S_arr[i-1]
        Qprev = Q_arr[i-1]
        # 计算当前行的S和Q
        S = Sprev + peti_arr[i] * np.exp(-1 * Sprev / smax_val) - precip_arr[i] + Qprev
        Q = precip_arr[i] + S
        S_arr[i] = S
        Q_arr[i] = Q
    return S_arr, Q_arr

# 加载数据
data = {'peti': {0: 0.1960418075323104, 1: 0.5796640515327454, 2: 0.737823486328125, 3: 0.222676545381546, 4: 0.8804306983947754}, 'tas': {0: 281.0088195800781, 1: 277.112060546875, 2: 273.7044372558594, 3: 277.48309326171875, 4: 279.4878845214844}, 'precip': {0: 0.0, 1: 0.0, 2: 1.5046296539367177e-05, 3: 0.0002500000118743, 4: 4.6296295295178425e-06}, 'year': {0: 2008, 1: 2008, 2: 2008, 3: 2008, 4: 2008}, 'row_id': {0: 0, 1: 1, 2: 2, 3: 3, 4: 4}, 'S': {0: 90.9, 1: "nan", 2: "nan", 3: "nan", 4: "nan"}, 'Q': {0: 0.0, 1: "nan", 2: "nan", 3: "nan", 4: "nan"}}
df = pd.DataFrame.from_dict(data)
smax_val = 100

start_time = time.time()
# 提取numpy数组（numba对numpy数组处理效率更高）
peti_arr = df['peti'].values
precip_arr = df['precip'].values
# 调用编译后的函数计算
S_arr, Q_arr = compute_s_q(peti_arr, precip_arr, smax_val, df['S'].iloc[0], df['Q'].iloc[0])
# 将结果写回DataFrame
df['S'] = S_arr
df['Q'] = Q_arr

print("--- %s seconds ---" % (time.time() - start_time))

为什么其他方法不适用？ #

• df.shift()：只能获取前一行的原始静态值，无法跟踪递推过程中每一步的动态计算结果。
• rolling.apply()：基于固定窗口的原始数据计算，无法处理依赖前一步实时输出的递推逻辑。
• 普通向量化：要求计算可并行执行，但你的场景下每一步输入依赖上一步输出，完全无法并行。

性能对比 #

针对3000行的数据，原始Python循环可能需要数百毫秒到几秒，而Numba编译后的循环仅需几毫秒，速度提升可达100倍以上。

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