如何高效批量处理Python Polars LazyFrame以写入数据库？

问题描述 #

需延迟加载大型Parquet文件，因数据库单次写入行数有限，需批量处理。由于LazyFrame无write_database方法，通过slice截取批次后collect为DataFrame，再调用DataFrame.write_database写入。但部分数据集下后续collect操作耗时剧增且内存占用上升：两列均为字符串的5000万行数据集，collect速度稳定；而含i64列的1亿行数据集，首次collect耗时约0.02秒，后续需10秒以上，且内存占用上升。各批次预估大小均为7-8MB。简化代码如下：

lf = pl.scan_parquet(file)

n_batch = 100_000
row = 0

while row < 500_000:
    lf_batch = lf.slice(row, n_batch)
    df_batch = lf_batch.collect()
    df_batch.write_database(…)
    row += n_batch

优化方案 #

• 复用执行计划，避免重复解析：每次调用slice都会重新生成查询计划，改用offset+limit组合，让Polars复用原LazyFrame的扫描计划，减少重复解析开销：

  lf = pl.scan_parquet(file)
  n_batch = 100_000
  total_rows = lf.select(pl.count()).collect().item()
  current_offset = 0
  
  while current_offset < total_rows:
      batch_df = lf.offset(current_offset).limit(n_batch).collect()
      batch_df.write_database(…)
      current_offset += n_batch
  

• 开启Parquet扫描缓存：创建LazyFrame时启用缓存，存储文件索引和统计信息，避免后续批次重复读取元数据：

  lf = pl.scan_parquet(file, cache=True)
  

• 明确指定数值列类型：针对含i64列的数据集，提前指定列类型，避免Polars重复执行类型推断：

  lf = pl.scan_parquet(
      file,
      dtype={"target_i64_col": pl.Int64}
  )
  

• 启用流式收集模式：在collect时添加streaming=True参数，让Polars以流方式处理数据，降低内存占用并提升分批效率：

  batch_df = lf.offset(current_offset).limit(n_batch).collect(streaming=True)
  

• 优化数据库写入参数：检查write_database是否支持批量写入参数（如batch_size），减少数据库连接和事务的重复开销，同时确保使用if_exists="append"模式避免重复建表。

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