← 中文资源

简体中文版本

Dask 分区并行聚合

用 Dask 分区 DataFrame 扩展 pandas 聚合:调度器并行执行各分区任务,分组聚合树形归并,map_partitions 自定义逐区函数。

阅读英文原文

示例说明

Dask 把大表切成分区,惰性构建任务图后并行执行:groupby 聚合被拆成各分区的局部聚合再树形合并,多核资源用满;map_partitions 则给每块应用任意 pandas 函数——比如取每区最大行——结果仍是分区化 DataFrame。compute 触发实际执行,其余时间都是纯声明。

技术要点

  • npartitions 控制并行任务粒度
  • 惰性任务图延迟到 compute 执行
  • meta 参数声明自定义分区函数的输出形态

代码解读

基准表十万行切成八分区;第一个查询按小时聚合均值、极值与样本数,Dask 自动做 map 侧局部聚合与 reduce 合并;第二个查询用 map_partitions 取每分区峰值行,meta 传空表头声明结构,compute 后 nlargest 取全局前三。

使用建议

分区数按核数与内存平衡,过多会调度开销占主导。meta 缺失会导致推断报错,自定义函数务必声明。超集群规模改用分布式调度器。

代码示例

import dask.dataframe as dd
import pandas as pd

base = pd.DataFrame({
    'sensor': [f's{i % 8}' for i in range(100_000)],
    'reading': [round(20 + (i % 37) * 0.5, 2) for i in range(100_000)],
    'hour': [(i // 3600) % 24 for i in range(100_000)],
})

ddf = dd.from_pandas(base, npartitions=8)
print('partition count:', ddf.npartitions)


def peak_of(partition: pd.DataFrame) -> pd.DataFrame:
    idx = partition['reading'].idxmax()
    return partition.loc[[idx]]


daily = (
    ddf
    .groupby('hour')
    .agg(
        mean_reading=('reading', 'mean'),
        max_reading=('reading', 'max'),
        samples=('reading', 'count'),
    )
    .compute()
)

peaks = (
    ddf
    .map_partitions(peak_of, meta=base.head(0))
    .compute()
    .nlargest(3, 'reading')
)

print('hours reported:', len(daily))
print('global top readings:')
print(peaks[['sensor', 'reading']].to_string(index=False))