如何在Dask中按列值分组数据并批量导出为CSV文件

你之前使用的groupby-apply范式存在天然局限：该操作需要将同一分组的所有数据拉取到单个worker节点的内存中处理，当分组基数大或单分组数据量超过内存时极易出现OOM问题，且确实只能调用Pandas接口，扩展性不足。以下两种原生Dask实现方案均支持超大规模数据集的分布式处理：

方案1：使用dask.dataframe.to_csv的partition_on参数（优先推荐） #

这是Dask官方原生支持的按字段拆分输出能力，全程不需要将单分组全量数据加载到同一节点，完全分布式执行，适配TB级以上数据集。

import dask.dataframe as dd

# 读取源数据，支持csv/parquet等任意Dask兼容的存储格式
ddf = dd.read_csv("your_source_data/*.csv", blocksize="100MB")

# 自定义输出文件名规则，入参为分组列的取值，返回对应输出文件名
def generate_filename(group_val):
    return f"output_{group_val}.csv"

# 直接导出，指定分组列即可自动拆分文件
ddf.to_csv(
    "./output_directory/",
    partition_on="your_group_column", # 替换为实际需要分组的列名
    name_function=generate_filename,
    index=False,
    encoding="utf-8"
)

• 该方法默认会在输出目录下生成分组列=取值格式的层级子目录，不需要层级的话可以直接在name_function的返回值中指定平级路径
• 支持高基数分组列，即使分组取值过万也可以正常运行，不需要提前做全量shuffle

方案2：按分组值过滤分批导出（适合需要自定义预处理的场景） #

如果每个分组导出前需要做自定义清洗、转换操作，可以先取分组列的唯一值列表，再逐个过滤导出，过滤操作本身也是分布式执行的：

import dask.dataframe as dd

ddf = dd.read_csv("your_source_data/*.csv", blocksize="100MB")
group_col = "your_group_column"

# 仅拉取分组列的唯一值列表到本地，数据量极小不会有性能问题
unique_group_vals = ddf[group_col].unique().compute()

# 遍历每个分组值过滤后导出
for val in unique_group_vals:
    group_df = ddf[ddf[group_col] == val]
    group_df.to_csv(
        f"./output_directory/group_{val}.csv",
        single_file=True, # 单分组体量不超过worker内存时可加该参数，让每个分组只输出1个CSV文件，超内存则移除该参数即可
        index=False,
        encoding="utf-8"
    )

• 如果分组取值数量超过10万，可将遍历逻辑改成Dask延迟任务批量提交，避免本地循环的性能开销

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