Dask的map_partitions为何比逐分区循环占用更多内存？

持有一份按车辆ID建立索引、按时间戳排序的车辆位置数据Parquet文件，需求为读取文件后对每个分区执行非聚合类计算，再将输出结果写入体量与原文件相近的新Parquet文件。
最初基于对Dask的认知使用map_partitions编写实现，原本预期该方法会逐分区执行操作、顺序将计算结果写入磁盘，可最小化内存占用，但实际运行时内存占用超出可用上限；改为手动循环逐分区执行计算、追加写入新Parquet文件后，程序可在内存限制内稳定运行。

待解答核心问题：

• 最初使用map_partitions的写法是否存在错误
• 如果写法无误，为什么该方法内存占用远高于手动循环实现
• 实现该需求的规范、最高效方式是什么

原始高内存占用实现（map_partitions版本） #

ddf = dd.read_parquet(input_file)
meta_dict = ddf.dtypes.to_dict()

(
    ddf
    .map_partitions(my_function, meta = meta_dict)
    .to_parquet(
        output_file,
        append = False,
        overwrite = True,
        engine = 'fastparquet'
    )
)

低内存占用手动循环实现 #

ddf = dd.read_parquet(input_file)

for partition in range(0, ddf.npartitions, 1):
    partition_df = ddf.partitions[partition]
    (
        my_function(partition_df)
        .to_parquet(
            output_file,
            append = True,
            overwrite = False,
            engine = 'fastparquet'
        )
    )

• 输入Parquet文件总大小约5GB，共拆分为11个分区，单分区最大体积900MB
• 数据按车辆ID设置索引与分区边界（divisions），所有车辆分组操作不需要跨分区完成
• 测试设备为16GB RAM、19GB交换分区的笔记本：原始map_partitions实现会占满内存与交换分区触发OOM，循环实现可以完全在RAM中运行。

1. map_partitions 写法不存在语法错误 #

你对map_partitions逐分区执行计算的基础逻辑认知正确，问题不在API调用语法层面，而在Dask默认的调度与写入机制。

2. 内存占用差异的核心原因 #

Dask默认多线程调度器执行完整任务图时，不会严格遵循「读取单个分区→计算单个分区→写入单个分区→释放内存」的串行流程：

• 调度器为了最大化计算吞吐量，会提前预读后续待处理分区到内存，默认配置下会同时加载多个分区的待处理数据。你的场景下单分区最大900MB，同时加载3-4个分区加上计算过程生成的临时对象，很容易突破16GB内存上限。
• 全局调用to_parquet时，Dask会在写入前执行额外的元数据校验、分区对齐检查，部分场景下还会缓存已计算完成的分区结果用于元数据生成，进一步抬高内存占用。
• 手动循环实现本质是强制将完整任务拆分为完全串行的独立子任务：每轮循环仅触发单个分区的读、算、写全流程，计算完成后该分区关联的所有对象会被垃圾回收机制直接释放，全程内存仅保留1个分区的数据，因此内存占用极低。

3. 规范高效的实现方案 #

无需手动编写循环，只需为Dask任务添加串行执行的配置约束，即可同时兼顾map_partitions的API简洁性和手动循环的低内存特性，推荐实现如下：

import dask

# 配置串行调度，关闭多任务并行预读
dask.config.set(scheduler='synchronous', optimize_graph=True)

ddf = dd.read_parquet(
    input_file,
    split_row_groups=False, # 严格映射Parquet原生行组为Dask分区，避免额外拆分产生开销
    engine='fastparquet'
)

(
    ddf
    .map_partitions(my_function, meta=ddf.dtypes.to_dict())
    .to_parquet(
        output_file,
        overwrite=True,
        engine='fastparquet',
        write_metadata_file=True, # 自动生成全局统一元数据，规避手动append可能引发的元数据损坏、分区读取顺序错乱问题
    )
)

该实现相比手动循环的优势是会自动生成标准Parquet元数据，稳定性更高。如果my_function为CPU密集型逻辑，也可替换为单worker的进程池调度，同样可以避免多分区同时加载导致的内存溢出问题，执行效率不会低于手动循环。

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