多进程下Polars sink_csv内存溢出问题及解决方案咨询

一、sink_csv 内置内存控制参数 #

Polars 的 sink_csv 提供了直接控制单进程内存占用的参数，可优先调整：

• batch_size：控制写入CSV的批次行数，默认值为 1024*1024。减小该值可降低每个批次在内存中的数据体积，比如设置 batch_size=100_000，让进程每次仅处理小批量数据后写入磁盘。
• compression：若启用压缩（如 gzip），压缩过程会额外消耗内存。内存紧张时可临时关闭压缩（compression=None），或切换到内存开销更低的压缩算法（如 lz4）。
• 聚合阶段前置优化：在合并 LazyFrames 并执行 sum/mean/min/max 等聚合操作时，全程保持 Lazy 模式，不要提前调用 collect()。Polars 会自动优化执行计划，避免全量数据加载到内存；若需强制流式聚合，可在相关操作中指定 streaming=True。

二、跨进程内存感知的实现方式 #

Polars 本身没有内置跨进程内存感知机制（进程内存空间隔离），但可通过外部工具实现动态调整：

• 借助 psutil 库在调用 sink_csv 前检查系统剩余内存，根据阈值动态调整 batch_size：

import psutil

def get_available_memory_gb():
    return psutil.virtual_memory().available / (1024**3)

# 按需调整批次大小
available_gb = get_available_memory_gb()
batch_size = 50_000 if available_gb < 20 else 200_000

lazy_df.sink_csv("output.csv", batch_size=batch_size)

• 极端情况下，可在内存不足时让进程短暂休眠（time.sleep()），等待其他进程释放内存后再继续写入。

三、最优实现方案 #

结合你的多进程+聚合+sink_csv 场景，最优方案是全链路流式处理，从源头避免内存堆积：

1. 坚持 Lazy 模式到底：
   合并 LazyFrames、执行聚合操作后，直接调用 sink_csv，全程不调用 collect()。Polars 会自动分批次处理数据，无需将全量聚合结果加载到内存。
2. 动态批次调整：
   配合 psutil 实现的内存检查逻辑，让每个进程根据系统实时内存状态调整写入批次，避免多个进程同时占用大量内存。
3. 系统层面兜底：
   若机器支持，可开启内存交换（swap）作为应急方案（注意会降低写入性能）；也可通过系统工具限制每个进程的内存配额（如 Linux 下的 cgroups），强制控制单进程内存上限。

针对聚合场景的额外提醒 #

合并 LazyFrames 时，确保使用 how="vertical"（默认值），Polars 可保持流式处理的优化空间；避免在聚合完成后调用 collect() 将数据转为内存中的 DataFrame，否则会导致全量数据驻留内存，写入 CSV 时内存中同时存在两份数据，大幅提升 OOM 风险。

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