如何用Polars高效连接大分区Parquet数据集并写入PostgreSQL？

针对你遇到的内存不足、分区读取效率低的问题，以下是几种基于Polars的高效处理方案：

一、启用Polars Streaming模式规避内存瓶颈 #

Polars的Streaming模式专为超大数据集设计，它会将数据拆分成多个批次逐块处理，无需一次性加载全部数据到内存，完美解决无法调用.collect()的问题，同时还能自动优化分区读取逻辑，减少重复读取S3分区的开销。

修改你的连接代码，添加streaming=True参数即可开启流式处理：

import polars as pl

df_left = pl.scan_parquet(
    "s3://my-bucket/left/**/*.parquet",
    hive_partitioning=True,
)
df_right = pl.scan_parquet(
    "s3://my-bucket/right/**/*.parquet",
    hive_partitioning=True,
)

# 开启流式内连接
df_joined = df_left.join(
    df_right,
    on=["category_id", "label_id"],
    how="inner",
    streaming=True
)

二、流式批量写入PostgreSQL #

Polars的write_database方法支持直接从延迟执行的LazyFrame写入数据库，配合streaming=True和batch_size参数，就能实现逐批次流式写入，无需将全量数据加载到内存。

示例代码：

from sqlalchemy import create_engine

# 创建PostgreSQL连接引擎
engine = create_engine("postgresql://<用户名>:<密码>@<主机>:<端口>/<数据库名>")

# 流式写入数据库
df_joined.write_database(
    table_name="joined_result",
    connection=engine,
    if_exists="replace",  # 根据需求选择"replace"或"append"
    batch_size=10_000,  # 调整批次大小（如1万~10万），平衡内存占用与写入效率
    streaming=True
)

三、优化分区读取效率：批量处理多category_id分区 #

如果你不想依赖Streaming模式，可以通过批量处理多个category_id分区来减少S3请求次数，降低单分区读取的开销：

1. 先获取所有唯一的category_id（仅需加载少量数据，内存压力小）
2. 将category_id分组，每组处理多个分区，减少S3连接次数

示例代码：

# 获取所有唯一的category_id
category_ids = df_left.select(pl.col("category_id")).unique().collect()["category_id"].to_list()

# 批量处理，每10个category_id为一组（可根据实际调整）
batch_size_cats = 10
for i in range(0, len(category_ids), batch_size_cats):
    current_cats = category_ids[i:i+batch_size_cats]
    # 过滤左右数据集到当前批次的category_id
    left_batch = df_left.filter(pl.col("category_id").is_in(current_cats))
    right_batch = df_right.filter(pl.col("category_id").is_in(current_cats))
    # 连接并追加写入数据库
    joined_batch = left_batch.join(right_batch, on=["category_id", "label_id"], how="inner")
    joined_batch.write_database(
        table_name="joined_result",
        connection=engine,
        if_exists="append",
        batch_size=10_000
    )

关键注意事项 #

• 确保使用最新版本的Polars（>=0.19.0），Streaming模式对join的支持在后续版本中持续优化，稳定性和效率更高。
• 调整batch_size参数时，需根据你的内存容量和数据库写入性能平衡：批次太小会增加数据库请求次数，太大则可能导致内存压力。

内容的提问来源于stack exchange，提问作者Joost Döbken