Polars LazyFrame大文件处理中如何保存中间结果，确保计算失败时已处理结果不丢失？

我完全懂你现在的糟心处境——用Polars流式处理大CSV时，中间的Python映射逻辑一报错，之前忙活半天处理的批次全白搭，输出文件空得离谱对吧？结合你的代码和场景，咱们先拆解问题根源，再给出落地的解决办法：

为什么现有方案会丢失数据？ #

你当前用map_batches+sink_csv的组合，哪怕开了streamable=True，但因为你的处理逻辑里用到了map_elements（属于Polars标记的OPAQUE_PYTHON节点），整个处理 pipeline 会被视为一个整体。sink_csv默认要等所有批次处理完成后才会把内容刷入磁盘，一旦中间某个批次报错，之前处理的内容都还在内存缓冲区里，没机会写入文件，自然就丢失了。调整batch_size或streaming_chunk_size也没用，因为核心问题是写入时机和处理流程的绑定关系。

核心解决思路：批次处理后立刻追加写入 #

要改成每个批次处理成功就立刻追加到输出文件，彻底把“处理”和“写入”的时机对齐，哪怕下一个批次失败，之前的成果也能保留。同时还要处理mapper的单点失败，避免单个行报错导致整个批次报废。

第一步：重构批次处理与写入逻辑 #

放弃依赖map_batches+sink_csv的自动流程，改用手动迭代流式批次，处理一个就写入一个：

import pathlib
import polars as pl

input_path = pathlib.Path("input.csv")
output_path = pathlib.Path("output.csv")
schema_as_dict = ... # 你的Polars schema字典
id_col_name = "id"
any_value_column_is_null = ... # 你的空值过滤表达式
generated_data_col_name = "generated_data"
mapper = ... # 你的昂贵计算函数
process_unprocessed_rows_in_batch = ... # 保留你原有的批次处理函数

# 初始化输出文件：如果不存在，先写入表头
if not output_path.exists():
    # 用schema生成空DataFrame写入表头
    pl.DataFrame(schema=schema_as_dict).write_csv(output_path, include_header=True)

# 流式迭代输入的每个批次
for batch in pl.scan_csv(input_path, schema=schema_as_dict).streaming_batches():
    try:
        # 处理当前批次
        processed_batch = process_unprocessed_rows_in_batch(batch)
        # 追加写入输出文件，跳过表头
        processed_batch.write_csv(
            output_path,
            include_header=False,
            mode="append",
            maintain_order=False
        )
        print(f"批次处理完成并写入")
    except Exception as e:
        print(f"批次处理失败，已保存之前的结果: {str(e)}")
        # 可选择记录失败批次的ID，方便后续排查
        break

第二步：修复mapper的单点失败问题 #

你的mapper偶尔会失败，如果不处理，单个行的报错会导致整个批次崩溃。我们可以从两方面优化：

1. 给mapper加异常捕获，返回符合schema的空值，不影响整个批次
2. 在map_elements里开启null_on_error=True（Polars 0.19+支持），单个元素失败时返回null而非崩溃

修改后的mapper和map_elements调用：

def mapper(row_id):
    try:
        # 你的昂贵计算逻辑
        # 返回符合schema_as_dict的字典
        return {"col1": computed_val1, "col2": computed_val2, ...}
    except Exception as e:
        print(f"处理ID {row_id}失败: {str(e)}")
        # 返回全空的字典，匹配schema
        return {col: None for col in schema_as_dict.keys()}

# 在process_unprocessed_rows_in_batch的map_elements里加null_on_error
pl.col(id_col_name).map_elements(
    mapper,
    pl.Struct(schema_as_dict),
    null_on_error=True  # 单个元素失败不影响整个批次
).alias(generated_data_col_name)

可选：断点续传优化 #

如果处理到一半崩溃，下次启动可以跳过已经处理过的行，避免重复劳动。可以通过读取输出文件的已处理ID来过滤输入：

# 读取已处理的ID集合
processed_ids = set()
if output_path.exists():
    # 流式读取输出文件的ID列，避免加载全量数据
    processed_ids = set(
        pl.scan_csv(output_path)
        .select(id_col_name)
        .collect(streaming=True)[id_col_name]
        .to_list()
    )

# 迭代输入时过滤已处理的行
for batch in pl.scan_csv(input_path, schema=schema_as_dict)\
               .filter(~pl.col(id_col_name).is_in(processed_ids))\
               .streaming_batches():
    # 后续处理逻辑同上
    pass

关键注意事项 #

• 表头一致性：第一次写入表头时，必须确保和后续追加的列顺序完全一致，否则会出现列错位的问题
• 写入性能：开启maintain_order=False可以大幅提升追加写入的速度，如果你不需要严格保持输入顺序的话
• 重试机制：如果mapper的失败是临时的（比如网络波动），可以用tenacity库给mapper加自动重试逻辑，减少失败概率

备注：内容来源于stack exchange，提问作者JRX