能否对依赖列历史值的Pandas DataFrame计算逻辑做向量化优化？

问题背景 #

当前通过循环遍历DataFrame计算Adj P*Q和Adj P两列，计算逻辑存在递推依赖：

• Adj P*Q（对应描述中的列F）依赖上一行的Adj P（对应描述中的列G）值（当Flag=-1时）
• Adj P（列G）依赖当前行的Adj P*Q值（当Flag=1时）

原始循环代码如下：

import pandas as pd
import numpy as np  # 补充原代码缺失的numpy导入

data = [[1,2131000,2131000,77.13],
[1,269000,2400000,79.25],
[1,1340000,3740000,81],
[1,268000,4008000,83.75],
[-1,1073000,2935000,85],
[1,269000,3204000,75]]

df = pd.DataFrame(data,columns=['Flag','Q','Cumsum Q','P'])

df['P*Q'] = df['P']*df['Q']
df.loc[0,'Adj P*Q'] = df.loc[0, 'P*Q']
df.loc[0,'Adj P'] = df.loc[0, 'P']

for index, row in df.iloc[1:,].iterrows():
        df.loc[index,'Adj P*Q'] = np.where(df.loc[index,'Flag'] == 1, df.loc[index-1,'Adj P*Q'] + df.loc[index,'P*Q'] * df.loc[index,'Flag'], df.loc[index-1,'Adj P'] * df.loc[index,'Cumsum Q'])
        df.loc[index,'Adj P'] = np.where(df.loc[index,'Flag'] == 1, df.loc[index,'Adj P*Q'] / df.loc[index,'Cumsum Q'],  df.loc[index-1,'Adj P'])

优化方案 #

这类带状态依赖的递推计算无法直接用常规Pandas向量化函数实现，但可以通过以下两种方式替代低效的iterrows循环：

方案1：用Numba编译加速循环 #

Numba可将Python循环编译为机器码，大幅提升计算速度，完美适配这类递推逻辑：

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

data = [[1,2131000,2131000,77.13],
[1,269000,2400000,79.25],
[1,1340000,3740000,81],
[1,268000,4008000,83.75],
[-1,1073000,2935000,85],
[1,269000,3204000,75]]

df = pd.DataFrame(data,columns=['Flag','Q','Cumsum Q','P'])
df['P*Q'] = df['P'] * df['Q']

# 提取数组用于Numba计算
flags = df['Flag'].values
p_q = df['P*Q'].values
cumsum_q = df['Cumsum Q'].values
n = len(df)

adj_pq = np.zeros(n, dtype=np.float64)
adj_p = np.zeros(n, dtype=np.float64)

# 初始化第一行
adj_pq[0] = p_q[0]
adj_p[0] = df.loc[0, 'P']

@jit(nopython=True)
def compute(flags, p_q, cumsum_q, adj_pq, adj_p):
    for i in range(1, n):
        if flags[i] == 1:
            adj_pq[i] = adj_pq[i-1] + p_q[i]
            adj_p[i] = adj_pq[i] / cumsum_q[i]
        else:
            adj_pq[i] = adj_p[i-1] * cumsum_q[i]
            adj_p[i] = adj_p[i-1]
    return adj_pq, adj_p

# 执行计算并赋值回DataFrame
df['Adj P*Q'], df['Adj P'] = compute(flags, p_q, cumsum_q, adj_pq, adj_p)

方案2：按Flag分段向量化计算 #

观察逻辑可发现，连续Flag=1的行可分组计算累加值，Flag=-1的行作为分段节点重置状态：

import pandas as pd
import numpy as np

data = [[1,2131000,2131000,77.13],
[1,269000,2400000,79.25],
[1,1340000,3740000,81],
[1,268000,4008000,83.75],
[-1,1073000,2935000,85],
[1,269000,3204000,75]]

df = pd.DataFrame(data,columns=['Flag','Q','Cumsum Q','P'])
df['P*Q'] = df['P'] * df['Q']

# 标记分段：每次Flag=-1时开启新分段
df['segment'] = (df['Flag'] == -1).cumsum()

# 初始化第一行
df.loc[0, 'Adj P*Q'] = df.loc[0, 'P*Q']
df.loc[0, 'Adj P'] = df.loc[0, 'P']

# 遍历每个分段（从第1段开始）
for seg in df['segment'].unique()[1:]:
    seg_mask = df['segment'] == seg
    first_idx = df[seg_mask].index[0]
    # 获取上一段最后一行的Adj P值
    prev_adj_p = df.loc[first_idx - 1, 'Adj P']
    
    # 处理分段第一行（Flag=-1）
    df.loc[first_idx, 'Adj P*Q'] = prev_adj_p * df.loc[first_idx, 'Cumsum Q']
    df.loc[first_idx, 'Adj P'] = prev_adj_p
    
    # 处理分段内后续的Flag=1行
    seg_rest_mask = seg_mask & (df.index > first_idx)
    if not seg_rest_mask.empty:
        # 计算P*Q的累加和，从分段第一行的Adj P*Q开始
        df.loc[seg_rest_mask, 'Adj P*Q'] = df.loc[first_idx, 'Adj P*Q'] + df.loc[seg_rest_mask, 'P*Q'].cumsum()
        # 计算对应Adj P
        df.loc[seg_rest_mask, 'Adj P'] = df.loc[seg_rest_mask, 'Adj P*Q'] / df.loc[seg_rest_mask, 'Cumsum Q']

# 删除辅助分段列
df.drop('segment', axis=1, inplace=True)

总结 #

• 由于计算逻辑存在状态递推依赖，无法完全用无状态的向量化函数实现，但上述两种方案可大幅提升效率：
  • Numba编译方案适合大数据量场景，速度提升最明显
  • 分段向量化方案无需额外依赖，逻辑更直观

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