如何让itertools.groupby兼容pd.NA？解决TypeError报错问题

在使用itertools.groupby处理包含pd.NA的序列时，会触发TypeError: boolean value of NA is ambiguous错误，最小复现代码如下：

import pandas as pd
import itertools

g = itertools.groupby([pd.NA,0])
next(g)
next(g)

问题原因 #

itertools.groupby的分组逻辑依赖相邻元素的相等性比较，但pd.NA参与任何比较运算都会返回NA，Python无法将NA转换为布尔值用于循环判断，因此触发报错。

解决方案 #

1. 预处理数据，替换pd.NA为唯一占位符

用一个自定义的唯一对象替换pd.NA，确保比较逻辑正常运行：

import pandas as pd
import itertools

data = [pd.NA, 0]
# 自定义唯一占位符，仅与自身相等
na_placeholder = object()
processed_data = [na_placeholder if pd.isna(x) else x for x in data]

g = itertools.groupby(processed_data)
print(next(g))  # 输出类似 (<object object at 0x10abc123>, <itertools._grouper object at 0x10def456>)
print(next(g))  # 输出 (0, <itertools._grouper object at 0x10ghi789>)

后续若需还原数据，将占位符替换回pd.NA即可。

2. 改用pandas原生的groupby方法

如果处理的是pandas Series，直接用Series自带的groupby更稳妥，它原生支持pd.NA分组：

import pandas as pd

s = pd.Series([pd.NA, 0])
# 获取分组对应的索引集合
group_indices = s.groupby(s).groups
print(group_indices)  # 输出 {pd.NA: [0], 0: [1]}

pandas的分组逻辑会将pd.NA视为独立的分组键，完全规避了itertools.groupby的比较问题，也更适配pandas数据结构的使用场景。

内容的提问来源于stack exchange，提问作者Michel de Ruiter