Pandas方法链中group by前后调用assign结果不同，如何统一输出？

问题根源 #

方案2出错的核心是你在assign里引用了原始的df.origin，而非方法链当前步骤的数据集列。经过query和groupby聚合后，数据集已经被过滤、压缩，行数远少于原始df，直接用df.origin会导致赋值时行数不匹配，结果错乱。而R的tidyverse中mutate默认操作的是当前管道内的数据，所以不会出现这类问题。

修正方案 #

在assign中用lambda x: ...引用方法链中当前的数据集，而非原始df。这样无论assign在链的哪个位置，都能操作当前上下文的数据。

修正后的方案2代码 #

(
 df
 .query("origin=='europe' & model_year==80")
 .groupby(['origin','cylinders'],dropna=False)
 .mpg
 .sum()
 .reset_index()
 .assign(origin=lambda x: x.origin.map({'europe':'Europe'}).fillna(x.origin))
)

额外说明 #

• 方案1中用df.origin能正常运行，是因为此时还没对数据集做过滤/聚合，原始df和当前链中的数据是同一套，但这种写法不够通用，推荐统一用lambda x:的方式，让代码在方法链任意位置都能稳定工作。
• 若要在分组后修改列，必须先通过reset_index()把分组索引转回普通列（你已经完成这一步），否则无法直接在assign中操作分组后的列。

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