如何优化大Polars LazyFrame分区Parquet写入以解决OOM问题？

核心问题分析 #

你当前的流程是先将Polars LazyFrame写入单一大临时文件，再用PyArrow扫描并分区写入，这种方式容易导致PyArrow读取大文件时内存过载，即使调整batch_size也难以缓解。以下是针对性的优化方案：

方案1：跳过临时文件，直接流式传递数据 #

直接将Polars LazyFrame的处理结果以Arrow批次的形式流式传给PyArrow，完全避免临时文件的磁盘IO和内存加载开销：

from loguru import logger
import polars as pl
import pyarrow.dataset
import pyarrow.parquet

# 加载并处理数据（全程保持LazyFrame，不要collect）
raw_data = pl.scan_parquet("raw_data.parquet")
# 示例处理步骤：raw_data = raw_data.filter(pl.col("value") > 0).with_columns(pl.col("date").cast(pl.Date))

# 将LazyFrame转为流式Arrow批次迭代器
arrow_batches = raw_data.to_arrow_batches()

# 用PyArrow流式写入分区Parquet
pyarrow.dataset.write_dataset(
    arrow_batches,
    "/path/to/data/",
    format="parquet",
    partitioning=["date"],
    partitioning_flavor="hive",
    max_partitions=1826,
    file_visitor=lambda f: logger.info(f"Wrote {f.path}"),
    existing_data_behavior="delete_matching",
    # 内存控制关键参数
    max_open_files=64,  # 限制同时打开的分区文件数，避免句柄和内存占用过高
    use_threads=False,  # 禁用多线程，避免并发处理导致内存堆积
    max_rows_per_file=1_000_000,  # 限制单个分区文件的行数，拆分大文件
    file_options=pyarrow.parquet.ParquetFileWriteOptions(
        write_statistics=False,  # 关闭Parquet统计信息写入，减少内存开销
        compression="snappy"  # 选择轻量压缩算法，平衡速度与内存
    )
)

方案2：优化临时文件流程（若必须落地临时数据） #

如果业务场景要求必须先落地临时数据，改为写入多个小文件而非单个大文件，降低PyArrow扫描时的内存压力：

import tempfile
from pathlib import Path
from loguru import logger
import polars as pl
import pyarrow.dataset

raw_data = pl.scan_parquet("raw_data.parquet")
# 保留你的数据处理步骤

with tempfile.TemporaryDirectory() as tmp_dir:
    # 将LazyFrame拆分为多个小文件写入临时目录
    raw_data.sink_parquet(
        Path(tmp_dir) / "part_*.parquet",
        row_group_size=100_000,  # 缩小行组大小，降低单文件内存占用
        compression="snappy",
        use_pyarrow=True  # 用PyArrow引擎写入，内存控制更稳定
    )
    
    # 扫描整个临时目录的小文件，而非单个大文件
    dataset = pyarrow.dataset.dataset(tmp_dir, format="parquet")
    scanner = dataset.scanner(batch_size=100_000)
    
    pyarrow.dataset.write_dataset(
        scanner,
        "/path/to/data/",
        format="parquet",
        partitioning=["date"],
        partitioning_flavor="hive",
        max_partitions=1826,
        file_visitor=lambda f: logger.info(f"Wrote {f.path}"),
        existing_data_behavior="delete_matching",
        max_open_files=64,
        max_rows_per_file=1_000_000
    )

方案3：排查处理流程中的内存隐患 #

确保数据处理全程使用Polars Lazy API，避免意外触发collect()导致全量数据加载到内存：

• 检查所有处理步骤，确保没有调用collect()、fetch()等强制计算的方法
• 用Polars的profile()分析处理流程的内存占用：

  # 生成处理流程的性能报告，查看内存峰值
  profile_result = raw_data.profile().collect()
  print(profile_result)
  

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