基于DataFrame分组转换：按当前行值的条件求和实现

需求说明 #

现有百万级行的DataFrame，字段包含article（商品）、shop（门店）、size（尺码）、units（数量），需要生成新列Result：

• 当当前行units不为NaN时，Result为该商品在对应门店下所有非NaN的units总和；
• 当当前行units为NaN时，Result直接设为NaN。
  要求避免逐行遍历，使用groupby+transform实现。

示例数据 #

import numpy as np
import pandas as pd

data = {
    "article": ("A", "A", "A", "A", "A", "A", "A"),
    "shop": ("x", "x", "y", "y", "z", "z", "z"),
    "size": ("M", "L", "M", "L", "M", "L", "S"),
    "units": (1, 2, np.nan, 4, np.nan, 6, 7)
}
df = pd.DataFrame(data)

期望Result列 #

[3, 3, np.nan, 4, np.nan, 13, 13]（对应1+2、1+2、空、4、空、6+7、6+7）

解决方案代码 #

# 按article和shop分组计算组内非NaN总和，再根据当前行units状态赋值
df["Result"] = df.groupby(["article", "shop"])["units"].transform(
    lambda g: g.sum(skipna=True)
).where(df["units"].notna(), np.nan)

代码解释 #

1. 分组求和：groupby(["article", "shop"])精准锁定「同一商品+同一门店」的分组范围，transform会为每组的每一行返回该组内所有非NaN的units总和，保证每行都能匹配对应分组的计算结果；
2. 条件过滤：利用where方法判断当前行units是否非空——非空则保留分组求和结果，为空则设为NaN，完全贴合需求逻辑。

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