如何用Dask DataFrame实现正确的分组自定义调和均值聚合？

我有一个超7000万行的大型数据集，需要对Dask DataFrame执行分组自定义聚合，聚合方式为调和均值。但当前的自定义聚合逻辑存在问题：Dask会先对每个分块做聚合，再对合并后的结果二次聚合，而分块计算的调和均值再合并，和直接对全量数据计算的结果不一致（比如[[1,3,5],[7,9]]的分块调和均值合并后，与[1,3,5,7,9]的整体调和均值不同）。

我的现有自定义聚合代码：

custom_hmean = dd.Aggregation('custom_hmean', lambda s: s.agg(hmean), lambda s0: s0.agg(hmean))

调和均值的本质是样本数量除以所有样本倒数的总和，公式为：
$$ H = \frac{n}{\sum_{i=1}^n \frac{1}{x_i}} $$

要在Dask中正确计算分组调和均值，不能直接分块计算调和均值再合并，而是要分三步聚合每个分组的样本总数和倒数之和，最后基于这两个值计算最终的调和均值：

1. 定义Dask聚合函数，分阶段处理统计值：

import dask.dataframe as dd
import pandas as pd

# 分块阶段：统计每个分组的样本数和倒数之和
def chunk_agg(s):
    return pd.Series({
        'count': s.count(),
        'sum_reciprocal': (1/s).sum()
    })

# 合并阶段：将所有分块的统计值按分组累加
def combine_agg(s):
    return pd.Series({
        'count': s['count'].sum(),
        'sum_reciprocal': s['sum_reciprocal'].sum()
    })

# 最终计算：用总样本数除以总倒数和得到调和均值
def final_agg(s):
    return s['count'] / s['sum_reciprocal']

# 定义自定义聚合规则
custom_hmean = dd.Aggregation(
    'custom_hmean',
    chunk=chunk_agg,
    combine=combine_agg,
    final=final_agg
)

2. 使用该聚合函数执行分组聚合：

# 假设你的Dask DataFrame为df，分组列为'group_col'，目标计算列为'value_col'
result = df.groupby('group_col')['value_col'].agg(custom_hmean)

核心说明 #

• 分块阶段：每个分块内仅统计可累加的基础值（样本数、倒数和），避免直接计算调和均值导致的误差
• 合并阶段：将各分块的统计值按分组累加，得到全局的样本总数和倒数总和
• 最终计算：基于全局统计值推导调和均值，结果与全量数据直接计算的结果完全一致

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