Pandas groupby与transform处理含NaN数据时运行异常怎么办？

问题背景 #

需要将组级统计值广播到分组内的每一行，当前使用的分组列存在多个NaN值。
此前使用以下代码实现需求，运行正常：

df.groupby(by="location")["id"].transform(lambda x: len(x))

旧版本运行时仅分组列为NaN的行不会得到对应统计值，选择id列作为操作列是为了避免对全列执行不必要的运算。

报错现象 #

处理新数据集时，只要分组变量包含NaN值就会抛出如下错误：

ValueError: Length mismatch: Expected axis has 316 elements, new values have 473 elements

问题排查 #

• 排除小分组影响：假设问题由大量仅含1行的分组导致，新增分组规模更大的列测试后，只要存在NaN值仍然触发相同错误
• 排除单列异常影响：数据框中所有可作为分组变量的列，存在NaN值时均报错，无NaN值时运行正常，排除单列出错的可能

解决方案 #

该问题为Pandas 1.2.5版本的固有缺陷，将Pandas升级到1.3.3及以上版本即可解决。 #

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