在np.where条件中使用pandas nullable integer dtype报错如何解决

报错原因 #

numpy的np.where执行时会先计算所有分支的输入表达式，不会因为外层条件过滤了空值就跳过第二个np.where的计算。第二个条件中的ydf['Z'] > 0在遇到<NA>时返回的结果还是pd.NA，numpy无法识别pandas的nullable NA类型，尝试将其转为布尔值时就会抛出boolean value of NA is ambiguous错误。

解决方案 #

方案1：使用pandas原生方法（推荐，兼容性最好） #

pandas的内置方法原生支持nullable数据类型，不需要额外适配：

# 逻辑：空值保留，大于0设为1，否则设为0，和你原代码逻辑一致
ydf['Z'] = ydf['Z'].where(pd.isna(ydf['Z']), (ydf['Z'] > 0).astype('Int32'))

如果你的需求是所有非空值直接替换为1，不需要判断大小，写法更简单：

ydf['Z'] = ydf['Z'].where(pd.isna(ydf['Z']), 1)

方案2：修改np.where的判断逻辑适配NA #

如果要继续使用np.where，需要在比较时提前处理空值的布尔结果，用pandas比较方法的fill_value参数将空值的比较结果设为False即可：

np.where(
    pd.isna(ydf['Z']), 
    pd.NA, 
    np.where(ydf['Z'].gt(0, fill_value=False), 1, 0)
)

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