如何用Pandas高效递归逐行计算并填充两列？百万行优化方案

给定含缺失值的P0列（记录指标初始值）和CHG列（记录百分比变化），需要生成P1列计算变化后的指标，同时用P1填充P0的缺失值。示例数据如下：

import pandas as pd
import numpy as np
df = pd.DataFrame()
df['P0'] = [50, np.nan, 60, np.nan]
df['CHG'] = [0, 0.5, -0.1, 0.2]

核心计算逻辑 #

1. 首行：P1[0] = P0[0] * (1 + CHG[0])，P0[0]保持原值
2. 后续行：
   • 若P0[i]缺失：P1[i] = 填充后的P0[i-1] * (1 + CHG[i])，并用P1[i]填充P0[i]
   • 若P0[i]非缺失：P1[i] = 填充后的P0[i-1] * (1 + CHG[i])，P0[i]保持原值

以下是两种适用于百万行数据的高效解决方案：

1. 向量化优化方案 #

利用pandas分组与累积乘积实现无循环计算，底层基于C优化，效率远高于纯Python循环。

import pandas as pd
import numpy as np

df = pd.DataFrame()
df['P0'] = [50, np.nan, 60, np.nan]
df['CHG'] = [0, 0.5, -0.1, 0.2]

# 标记P0非缺失位置为分段起点，生成分组ID
segment_markers = df['P0'].notna()
df['segment_id'] = segment_markers.cumsum()

# 计算每个分段内(1+CHG)的累积乘积
df['cum_prod'] = df.groupby('segment_id')['CHG'].apply(lambda x: (1 + x).cumprod())

# 映射每个分段的初始P0值
segment_starts = df.loc[segment_markers, ['segment_id', 'P0']].set_index('segment_id')['P0']
df['segment_start'] = df['segment_id'].map(segment_starts)

# 计算填充后的P0序列
df['shifted_cum_prod'] = df.groupby('segment_id')['cum_prod'].shift(1).fillna(1)
df['P0_filled'] = df['segment_start'] * df['shifted_cum_prod']

# 计算P1列
df['P1'] = df['P0_filled'].shift(1).fillna(df['segment_start'].iloc[0]) * (1 + df['CHG'])

# 更新原P0列并清理临时列
df['P0'] = df['P0_filled']
df.drop(['segment_id', 'cum_prod', 'segment_start', 'shifted_cum_prod', 'P0_filled'], axis=1, inplace=True)

运行结果：

2. Numba加速循环方案 #

针对复杂逻辑或超大规模数据，Numba可将Python循环编译为机器码，速度接近C语言，比纯Python循环快数十倍。

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

@jit(nopython=True)
def compute_p0_p1(p0, chg):
    n = len(p0)
    p0_filled = p0.copy()
    p1 = np.empty(n, dtype=np.float64)
    
    # 处理首行
    p1[0] = p0_filled[0] * (1 + chg[0])
    
    for i in range(1, n):
        # 计算当前P1
        p1[i] = p0_filled[i-1] * (1 + chg[i])
        # 填充缺失的P0
        if np.isnan(p0_filled[i]):
            p0_filled[i] = p1[i]
    
    return p0_filled, p1

# 转换为numpy数组（Numba对numpy优化效果最优）
p0_np = df['P0'].values.copy()
chg_np = df['CHG'].values.copy()

# 调用加速函数并更新DataFrame
filled_p0, p1_result = compute_p0_p1(p0_np, chg_np)
df['P0'] = filled_p0
df['P1'] = p1_result

方案对比 #

• 向量化方案：代码简洁，依赖pandas原生优化，无需额外库，适合常规场景，处理百万行数据效率充足。
• Numba方案：适配复杂逻辑场景，加速效果显著，但需安装Numba库，代码需符合其nopython模式语法要求。

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