python - 在处理 NaN 时如何快速对 Pandas groupby 对象求和？

Question

我有一个带有 key 和 value 列的 DataFrame 。 值有时不适用:

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。

this answer中的代码在我的机器上需要 35.7 秒:

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

这比理论上可能的速度要慢得多。内置的 Pandas SeriesGroupBy.sum 在我的机器上需要 6.31 秒:

df.groupby('key')['value'].sum()

但它不支持 NA 处理(请参阅 this GitHub issue )。

我可以编写哪些代码来获得与内置运算符相当的性能，同时仍然处理 NaN？

Answer 1

一种解决方法可能是用 Inf 替换 NaN:

df.fillna({'value': np.inf}).groupby('key')['value'].sum().replace(np.inf, np.nan)

更快的替代方案:

df['value'].fillna(np.inf).groupby(df['key']).sum().replace(np.inf, np.nan)

示例输出:

key
0        45208.0
1            NaN
2        62754.0
3        50001.0
4        51073.0
          ...   
99995    55102.0
99996    43048.0
99997    49497.0
99998    43301.0
99999        NaN
Name: value, Length: 100000, dtype: float64

计时(10m 行)。

# original sum
743 ms ± 81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Inf workaround
918 ms ± 70.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Inf workaround (alternative)
773 ms ± 60.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# custom apply with numpy
5.99 s ± 263 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)