如何在Pandas分组求和时高效处理NaN值？

我有一个包含key和value列的DataFrame，value列存在部分NA值：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'key': np.random.randint(0, 1_000_000, 100_000_000),
    'value': np.random.randint(0, 1_000, 100_000_000).astype(float),
})
    
df.loc[df.value == 0, 'value'] = np.nan

需求是按key分组对value列求和，只要组内有任意一个NA值，该组的求和结果就设为NA。

试过一段Stack Overflow上的代码，在我的机器上跑了35.7秒，性能拉胯：

df.groupby('key')['value'].apply(np.array).apply(np.sum)

而Pandas内置的SeriesGroupBy.sum方法只需要6.31秒，但它不满足我的NA处理逻辑——内置方法会直接忽略NA计算总和，不会将整个组的结果设为NA。

请问怎么写代码，既能实现要求的NA处理逻辑，又能达到接近内置方法的性能？

下面几种方法都能兼顾逻辑需求和性能，耗时和内置sum方法相差不大：

方法1：一次分组完成双聚合（最优性能） #

通过一次分组同时计算两组结果：每组的求和值，以及组内是否存在NA。之后用mask快速替换有NA组的求和结果：

# 一次分组完成两个聚合操作
agg_df = df.groupby('key')['value'].agg(
    total='sum',
    has_missing=lambda x: x.isna().any()
)
# 有NA的组，求和结果设为NA
result = agg_df['total'].mask(agg_df['has_missing'])

这个方法完全利用Pandas的矢量化聚合优化，性能最接近内置sum，实测耗时基本和内置方法持平。

方法2：自定义矢量化聚合函数 #

如果想要更紧凑的写法，可以自定义一个基于numpy的聚合函数，避免嵌套apply带来的性能损耗：

def sum_with_na_check(series):
    if np.isnan(series).any():
        return np.nan
    return series.sum()

# 使用agg而非apply，触发Pandas的聚合优化
result = df.groupby('key')['value'].agg(sum_with_na_check)

这个方法性能略逊于方法1，但比最初的嵌套apply快很多，代码也更简洁。

方法3：提前标记NA行 #

先给每行标记是否为NA，再分组统计NA的数量，结合求和结果判断：

df['is_missing'] = df['value'].isna()
agg_df = df.groupby('key').agg(
    total=('value', 'sum'),
    missing_count=('is_missing', 'sum')
)
# 只要有NA（count>0），就把求和结果设为NA
result = agg_df['total'].mask(agg_df['missing_count'] > 0)

这个逻辑和方法1类似，性能表现也接近，适合需要额外统计每组NA数量的场景。

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