Dask分组聚合中自定义带ddof参数的方差函数报错问题求助

我最近在使用Dask做分组聚合时遇到了一个棘手的问题：想要实现带自定义ddof参数的方差计算，同时和其他聚合操作一起执行，但总是触发RuntimeError，想请教各位大佬问题出在哪里，有没有更合适的解决方法？

背景情况 #

单独计算指定非默认ddof的分组方差是可以正常运行的：

df.groupby("a")["b"].var(ddof=2)

如果是和其他聚合操作搭配使用默认方差的话，这样写也没有问题：

df.groupby("a").agg(b_var = ("b", "var"), c_sum = ("c", "sum"))

尝试的自定义聚合函数 #

为了实现带自定义ddof的方差聚合，我参考Dask文档写了一个自定义聚合函数：

def var(ddof: int = 1) -> dd.Aggregation:
    import dask.dataframe as dd

    return dd.Aggregation(
        name="var",
        chunk=lambda s: (s.count(), s.sum(), (s.pow(2)).sum()),
        agg=lambda count, sum_, sum_sq: (count.sum(), sum_.sum(), sum_sq.sum()),
        finalize=lambda count, sum_, sum_sq: (sum_sq - (sum_ ** 2 / count)) / (count - ddof),
    )

报错情况 #

但当我执行下面的聚合代码时，就会抛出RuntimeError：

df.groupby("a").agg({"b": var(2)})

具体报错信息如下：

RuntimeError('Failed to generate metadata for DecomposableGroupbyAggregation(frame=df, arg={‘b’ <dask.dataframe.groupby.Aggregation object at 0x7fdfb8469910>}

我还尝试把代码里的s.pow(2)替换成s**2，但依然会报错。

完整测试脚本 #

import dask.dataframe as dd

data = {
    "a": [1, 1, 1, 1, 2, 2, 2],
    "b": range(7),
    "c": range(10, 3, -1),
}

df = dd.from_dict(data, 2)

def var(ddof: int = 1) -> dd.Aggregation:
    import dask.dataframe as dd

    return dd.Aggregation(
        name="var",
        chunk=lambda s: (s.count(), s.sum(), (s.pow(2)).sum()),
        agg=lambda count, sum_, sum_sq: (count.sum(), sum_.sum(), sum_sq.sum()),
        finalize=lambda count, sum_, sum_sq: (sum_sq - (sum_ ** 2 / count)) / (count - ddof),
    )

df.groupby("a").agg(b_var = ("b", "var"), c_sum = ("c", "sum"))  # <- 这行正常运行
df.groupby("a").agg(b_var = ("b", var(2)), c_sum = ("c", "sum"))  # <- 这行触发RuntimeError

备注：内容来源于stack exchange，提问作者FBruzzesi