Dask DataFrame分组中使用自定义聚合函数的几种常用方法 - 随笔

方案1: 使用 agg 结合Lambda函数

适用于简单聚合，直接使用内置函数与Lambda表达式组合。

import dask.dataframe as dd
import pandas as pd
from dask.dataframe import Aggregation

# 创建示例数据
pandas_df = pd.DataFrame({
    'group': ['A', 'A', 'B', 'B', 'C', 'C'],
    'value1': [1, 2, 3, 4, 5, 6],
    'value2': [10, 20, 30, 40, 50, 60]
})
dask_df = dd.from_pandas(pandas_df, npartitions=2)

# 定义自定义聚合函数
range_agg = Aggregation('range', lambda s: s.max() - s.min(), lambda s: s.max() - s.min())
sum_times_01_agg = Aggregation('sum_times_01', lambda s: s.sum() * 0.1, lambda s: s.sum() * 0.1)

# 分组聚合：使用预定义的聚合对象
result = dask_df.groupby('group').agg({
    'value1': ['mean', range_agg],  # 均值与极差
    'value2': [sum_times_01_agg]    # 和的10%
}).compute()

print(result)

该方法同时可以避免出现异常：raise ValueError(f"unknown aggregate {func}") ValueError: unknown aggregate lambda

输出

        mean range sum_times_01
group
A        1.5     0          0.3
B        3.5     0          0.7
C        5.5     0          1.1

---

方案2: 使用 apply 自定义聚合函数

适用于复杂聚合逻辑，需返回单个标量或多个结果。

import dask.dataframe as dd
import pandas as pd
from dask.dataframe import Aggregation

pandas_df = pd.DataFrame({
    'group': ['A', 'A', 'B', 'B', 'C', 'C'],
    'value1': [1, 2, 3, 4, 5, 6],
    'value2': [10, 20, 30, 40, 50, 60]
})
dask_df = dd.from_pandas(pandas_df, npartitions=2)

# 定义返回单行DataFrame的聚合函数
def custom_aggregation(df):
    return pd.Series({
        'v1_mean': df['value1'].mean(),
        'v1_range': df['value1'].max() - df['value1'].min(),
        'v2_total': df['value2'].sum()
    })

# 应用时需指定输出的元数据（meta）
result = dask_df.groupby('group').apply(
    custom_aggregation, 
    meta={'v1_mean': 'f8', 'v1_range': 'f8', 'v2_total': 'f8'}
).compute()

print(result)

输出

       v1_mean  v1_range  v2_total
group
A          1.5       1.0      30.0
B          3.5       1.0      70.0
C          5.5       1.0     110.0

---

方案3: 使用 namedagg（类似Pandas的命名聚合）

Dask支持类似Pandas的命名聚合方式，代码更清晰。

import dask.dataframe as dd
import pandas as pd
from dask.dataframe import Aggregation

# 创建示例数据
pandas_df = pd.DataFrame({
    'group': ['A', 'A', 'B', 'B', 'C', 'C'],
    'value1': [1, 2, 3, 4, 5, 6],
    'value2': [10, 20, 30, 40, 50, 60]
})
dask_df = dd.from_pandas(pandas_df, npartitions=2)

range_agg = Aggregation('range', lambda s: s.max() - s.min(), lambda s: s.max() - s.min())

# 使用 namedagg 语法
result = dask_df.groupby('group').agg(
    v1_mean=('value1', 'mean'),
    v1_range=('value1', range_agg),  # 自定义 Lambda
    v2_sum=('value2', 'sum')
).compute()

print(result)

输出

       v1_mean  v1_range  v2_sum
group
A          1.5         0      30
B          3.5         0      70
C          5.5         0     110

---

方案4: 使用 map_partitions 分块处理

在分区级别先预处理，再全局聚合，适合超大数据集。

import dask.dataframe as dd
import pandas as pd

# 创建示例数据
pandas_df = pd.DataFrame({
    'group': ['A', 'A', 'B', 'B', 'C', 'C'],
    'value1': [1, 2, 3, 4, 5, 6],
    'value2': [10, 20, 30, 40, 50, 60]
})
dask_df = dd.from_pandas(pandas_df, npartitions=2)

# 每个分区内预处理（如过滤、计算中间量）
def preprocess(partition):
    ## 1. 直接操作（可能发出警告）
    # partition['value1_squared'] = partition['value1'] ** 2
    
    ## 2. 创建分区数据的副本
    partition_copy = partition.copy()
    # 在副本上进行操作
    partition_copy['value1_squared'] = partition_copy['value1'] ** 2

    return partition_copy # partition

dask_df_processed = dask_df.map_partitions(preprocess)

# 全局聚合
result = dask_df_processed.groupby('group').agg({
    'value1_squared': 'sum',
    'value2': 'mean'
}).compute()

print(result)

输出

       value1_squared  value2
group
A                   5    15.0
B                  25    35.0
C                  61    55.0

---

方案5: 混合内置与自定义函数的多列聚合

同时对多列应用不同聚合操作。

import dask.dataframe as dd
import pandas as pd
from dask import delayed

# 创建示例数据
pandas_df = pd.DataFrame({
    'group': ['A', 'A', 'B', 'B', 'C', 'C'],
    'value1': [1, 2, 3, 4, 5, 6],
    'value2': [10, 20, 30, 40, 50, 60]
})
dask_df = dd.from_pandas(pandas_df, npartitions=2)

# 定义多列混合聚合逻辑
def multi_column_agg(df):
    return pd.DataFrame({
        'group': df['group'].iloc[0],
        'v1_sum': [df['value1'].sum()],
        'v2_ratio': [df['value2'].max() / df['value2'].min()]
    })

# 应用聚合，指定输出结构
result = dask_df.groupby('group').apply(
    multi_column_agg,
    meta={'group': 'object', 'v1_sum': 'f8', 'v2_ratio': 'f8'}
).compute()


'''
# 也可以 dask.delayed 包装整个分组聚合过程（不推荐，仅作演示）
@delayed
def delayed_groupby_apply(df, func, meta):
    return df.groupby('group').apply(func, meta=meta).compute()  # 注意：这里直接调用了 .compute()，实际上失去了延迟的意义

# 使用 dask.compute 执行延迟任务（但在这个例子中，由于 .compute() 已经在 delayed 函数内部被调用，这一步实际上是多余的）
# 正确的做法是直接调用 dask_df.groupby().apply().compute()
# result = compute(result_delayed)[0]  # 这行代码在这个例子中不会按预期工作，因为 .compute() 已经在 delayed 函数内部被调用

'''

print(result)

输出

        group  v1_sum  v2_ratio
group
A     0     A       3  2.000000
B     0     B       7  1.333333
C     0     C      11  1.200000

---

关键说明

1. 延迟执行：Dask操作默认惰性执行，需调用 .compute() 触发计算。
2. 元数据（meta）：使用 apply 时必须通过 meta 参数明确输出结构。
3. 性能优化：若数据能放入内存，直接使用Pandas更高效；Dask适用于分布式大数据场景。

实际应用中应根据需求选择合适方案，平衡代码可读性与计算效率。