如何高效为超大型DataFrame添加同SKU上周库存列？

需求说明 #

现有包含sku_id、week、stock三列的数值型DataFrame，需新增stock_week_minus_one列，存储对应SKU上一周的库存数据。原始数据未按sku_id或week排序，原使用iterrows()遍历的方法在180万行数据下运行极慢，需更高效的实现方式。

输入示例 #

sku_id  week  stock
0       1    10    100
1       2     1     50
2       2     3     10
3       2     2     30
4       1    11     50

期望输出 #

sku_id  week  stock  stock_week_minus_one
0       1    10    100                   NaN
1       2     1     50                   NaN
2       2     3     10                  30.0
3       2     2     30                  50.0
4       1    11     50                 100.0

低效原方法问题分析 #

原代码通过iterrows()遍历每一行，每次循环全表查询对应SKU的上一周库存，时间复杂度为O(n²)，在180万行数据下会产生大量重复计算，导致运行速度极慢。

高效优化方案 #

利用pandas的向量化分组+移位操作，将时间复杂度降至O(n log n)（主要来自排序步骤），大幅提升处理速度。

基础实现（适用于SKU周数连续的场景） #

import pandas as pd

# 1. 按sku_id和week排序，确保同SKU的周数据按顺序排列
df_sorted = df.sort_values(['sku_id', 'week'])
# 2. 按sku_id分组，对stock列执行移位操作，shift(1)取上一行的库存值
df_sorted['stock_week_minus_one'] = df_sorted.groupby('sku_id')['stock'].shift(1)
# 3. 恢复原数据的索引顺序（若无需保留原顺序可省略）
df = df_sorted.sort_index()

进阶实现（处理SKU周数不连续的场景） #

若部分SKU存在周数缺失（如存在week=1和week=3，但无week=2），需额外验证当前周是否为上一周的连续周，否则将对应值设为NaN：

import pandas as pd
import numpy as np

df_sorted = df.sort_values(['sku_id', 'week'])
# 执行移位操作
df_sorted['stock_week_minus_one'] = df_sorted.groupby('sku_id')['stock'].shift(1)
# 生成掩码：检查当前周是否等于上一周+1
mask = df_sorted.groupby('sku_id')['week'].shift(1) + 1 == df_sorted['week']
# 仅保留连续周的移位值，非连续周设为NaN
df_sorted['stock_week_minus_one'] = df_sorted['stock_week_minus_one'].where(mask, other=np.nan)
# 恢复原顺序
df = df_sorted.sort_index()

方案优势 #

• 基于pandas内置的向量化操作，避免了逐行遍历的低效计算
• 分组与移位操作均为pandas底层优化的C实现，处理180万行数据仅需数秒
• 代码简洁易维护，可灵活适配不同数据场景

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