如何用Polars将Python函数生成的数据流式写入Parquet文件？

我有一个Python脚本会对文件列表执行操作，生成的DataFrame约6GB，因此需要用Parquet压缩存储。问题在于数据是通过函数逐批生成的，不想先把所有数据存到内存里再拆分写入，希望能直接流式写入Parquet。我主要用Polars而非Pandas，求相关实现建议。

附上当前代码：

from time import sleep


def min_and_max(x: list[int]) -> tuple[int, int]:
    sleep(30)  # 模拟耗时操作
    return (min(x), max(x))


numbers = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]

for num in numbers:
    print(min_and_max(num))
    # 这里需要写入Parquet的逻辑
    # 把结果存在列表或其他数据结构里太占内存

方法1：分批追加写入Parquet #

Polars支持写入Parquet时使用mode="append"参数，你可以每次生成一小批数据就直接写入文件，无需将全量数据驻留内存。

修改后的代码示例：

import polars as pl
from time import sleep


def min_and_max(x: list[int]) -> tuple[int, int]:
    sleep(30)  # 模拟耗时操作
    return (min(x), max(x))


numbers = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
output_path = "results.parquet"

# 初始化：处理第一个批次并创建Parquet文件
first_num = numbers[0]
min_val, max_val = min_and_max(first_num)
df_first = pl.DataFrame({"min": [min_val], "max": [max_val]})
df_first.write_parquet(output_path)

# 处理剩余批次并追加写入
for num in numbers[1:]:
    min_val, max_val = min_and_max(num)
    df_batch = pl.DataFrame({"min": [min_val], "max": [max_val]})
    df_batch.write_parquet(output_path, mode="append")

注意：这种方式适合单条数据量极小的场景，但频繁追加可能会生成较多小文件块，后续读取时可通过pl.scan_parquet自动优化性能。

方法2：迭代器流式生成写入 #

如果希望更贴合流式处理的逻辑，可以把数据生成逻辑包装成迭代器，转换成Polars的LazyFrame后直接写入，全程内存占用极低：

import polars as pl
from time import sleep


def min_and_max(x: list[int]) -> tuple[int, int]:
    sleep(30)  # 模拟耗时操作
    return (min(x), max(x))


def data_generator(numbers_list):
    for num in numbers_list:
        min_val, max_val = min_and_max(num)
        yield {"min": min_val, "max": max_val}


numbers = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]

# 从迭代器创建流式LazyFrame
lf = pl.scan_batches(data_generator(numbers))

# 流式写入Parquet
lf.sink_parquet("results_streaming.parquet")

这种方式下，Polars会自动从迭代器逐批读取数据并写入文件，完全不需要加载全量数据到内存，适合处理大规模生成的数据。

额外优化建议 #

• 调整批次大小：如果单条数据太小，可以攒N个结果再生成批次写入，减少IO次数提升效率，比如每处理100条数据写入一次。
• 指定压缩算法：写入时可选择更高压缩比的算法（如zstd）进一步缩小文件体积：

  lf.sink_parquet("results_streaming.parquet", compression="zstd")
  

• 显式指定Schema：提前定义数据结构，避免Polars自动推断的开销：

  schema = pl.Schema({"min": pl.Int64, "max": pl.Int64})
  lf = pl.scan_batches(data_generator(numbers), schema=schema)
  

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