如何在Dask聚合函数中传参？解决全NaN列求和返回NaN的问题

为什么全NaN列求和默认返回0？ #

Pandas和Dask的sum()函数默认参数min_count=0，这个参数的作用是：当非NaN值的数量小于min_count时，返回NaN。默认设为0意味着哪怕没有任何非NaN值（全NaN），也会返回0。你想要的全NaN返回NaN的效果，其实需要设置min_count=1——只要没有至少1个非NaN值，就返回NaN。

Dask GroupBy聚合时的解决方法 #

你直接在aggregate()里加min_count参数没用，是因为传字符串'sum'时，Dask调用的是默认的sum函数，不会传递额外参数；用lambda报错则是因为Dask的分布式计算需要函数能被序列化，lambda在序列化时容易出问题。可以用以下两种方法解决：

方法1：用functools.partial包装sum函数

通过partial给sum绑定min_count=1参数，生成一个可被Dask序列化的新函数：

from functools import partial
import dask.dataframe as dd

# 包装sum函数，指定min_count=1
sum_with_nan = partial(sum, min_count=1)

# 执行聚合
result = ddf.groupby("code").aggregate({'rain': sum_with_nan}).compute()

方法2：定义自定义聚合函数

直接写一个明确调用sum(min_count=1)的函数，同样能被Dask正确序列化：

import dask.dataframe as dd

def nan_safe_sum(series):
    return series.sum(min_count=1)

# 执行聚合
result = ddf.groupby("code").aggregate({'rain': nan_safe_sum}).compute()

验证效果 #

用全NaN的测试数据验证：

import pandas as pd
import dask.dataframe as dd

# 创建测试数据
data = {'code': ['A', 'A', 'B', 'B'], 'rain': [None, None, None, None]}
df = pd.DataFrame(data)
ddf = dd.from_pandas(df, npartitions=2)

# 用上面的方法聚合，结果中A和B的rain列都会返回NaN，而不是0
print(result)

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