在DataFrameGroupBy.agg中，sum(min_count=1)是否与sum一样快速？

Question

我有一个包含几百万行的数据表，我通过某种方式对它进行分组，然后对几个列进行求和，代码如下：

df.groupby(…).agg({
  "foo": "sum",
  "bar": "sum",
  …})

这个过程相当快，我的电脑上大约需要10秒钟。

现在我需要一种求和的变体，它可以保留NaN值，比如说：

sum(min_count=1)

看起来要想用上面提到的分组和聚合方法实现这个功能，只能通过以下方式：

... agg({
  "foo": (lambda s: s.sum(min_count=1))
})

但是这样做的速度慢了10倍。

我尝试传递一些参数，但当我把一个字典传给聚合方法时，这些参数似乎没有被使用：

... agg({
  "foo": "sum",
  "bar": "sum",
}, min_count=1)

有没有办法在保持min_count=1的语义的同时，达到和简单求和一样的速度呢？

Answer 1

你可以尝试的一种方法是直接在自定义的聚合函数中使用numpy。

import numpy as np

def sum_with_min_count(s):
    return np.nansum(s.to_numpy())

result = df.groupby(...).agg({
    "foo": sum_with_min_count,
    "bar": sum_with_min_count,
    # Add other columns as needed
})

Answer 2

如果你在所有列上使用 sum 作为聚合函数，你可以这样做：

df.groupby(...).sum(min_count=1)

---

以这个数据框为例：

   a    b    c
0  1  NaN  NaN
1  1  NaN  NaN
2  2  1.0  3.0
3  2  2.0  4.0

这样做：

df.groupby("a").sum(min_count=1)

会返回：

     b    c
a          
1  NaN  NaN
2  3.0  7.0

---

如果你想对特定的列进行聚合，可以在分组后选择它们：

   a    b    c    d
0  1  NaN  NaN  NaN
1  1  NaN  NaN  NaN
2  2  1.0  3.0  0.0
3  2  2.0  4.0  0.0

比如选择 'b' 和 'c'：

df = df.groupby("a")[["b", "c"]].sum(min_count=1)

会返回：

     b    c
a          
1  NaN  NaN
2  3.0  7.0

Answer 3

我觉得问题的关键在于，你可以把：

df.groupby(...).agg({"col1": func0, "col2": func0})

替换成：

df.groupby(...)[["col1", "col2"]].func0()。

在你的例子中：

>>> df = pd.DataFrame({"x": list("aabbbaa"), "y": [1,1,None,None,None,None,None], "z": [None,None,3,None,None,None,None]})
>>> df.groupby("x")[["y", "z"]].sum(min_count=1)
     y    z
x
a  2.0  NaN
b  NaN  3.0

这让你可以充分利用 pd.DataFrame.sum() 的所有功能，具体可以参考 文档。