初始化1500万条记录的Polars DataFrame致机器死机，求优化方案

当前使用Psycopg3从AWS RDS Postgres拉取数据，初始化Polars DataFrame时，100万条记录运行正常，但1500万条记录会导致机器死机重启。Psycopg连接器可正常获取全量数据，问题仅出在Polars DataFrame初始化环节，尝试改用LazyFrame后问题依旧。原实现代码如下：

rds_conn = psycopg.connect(
        host=config.RDS_HOST_NAME,
        dbname=config.RDS_DB_NAME,
        user=config.RDS_DB_USER,
        password=config.RDS_DB_PASSWORD,
        port=config.RDS_PORT)
    cur = rds_conn.cursor(name="rds_cursor")
    cur.itersize = 100000
    cur.execute(sql_query)
    names = [x[0] for x in cur.description]
    rows = cur.fetchall()
    cur.close()
    df = pl.DataFrame(rows, schema=names, infer_schema_length=None)

优化方案 #

1. 分批次加载并增量合并 #

避免一次性将所有数据加载到内存，通过循环分批获取数据，逐步合并成最终DataFrame，控制单批次内存占用：

rds_conn = psycopg.connect(
    host=config.RDS_HOST_NAME,
    dbname=config.RDS_DB_NAME,
    user=config.RDS_DB_USER,
    password=config.RDS_DB_PASSWORD,
    port=config.RDS_PORT
)
cur = rds_conn.cursor(name="rds_cursor")
cur.itersize = 100000
cur.execute(sql_query)
names = [x[0] for x in cur.description]

# 初始化空DataFrame
df = pl.DataFrame(schema=names)

# 循环分批加载合并
while batch := cur.fetchmany(100000):
    batch_df = pl.DataFrame(batch, schema=names, infer_schema_length=None)
    df = df.vstack(batch_df)

cur.close()
rds_conn.close()

2. 使用Polars原生Postgres读取接口 #

Polars内置的pl.read_database()方法会自动优化数据加载流程，底层处理分批和内存管理，支持直接生成LazyFrame延迟计算，大幅降低内存压力：

# 生成LazyFrame（推荐，延迟计算仅在需要时执行）
lf = pl.read_database(
    sql_query,
    connection=f"postgresql://{config.RDS_DB_USER}:{config.RDS_DB_PASSWORD}@{config.RDS_HOST_NAME}:{config.RDS_PORT}/{config.RDS_DB_NAME}",
    engine="psycopg"
)

# 若需立即转换为DataFrame，执行：
# df = lf.collect()

3. 自定义紧凑数据类型 #

手动指定更节省内存的数据类型（如用pl.Int32替代pl.Int64、pl.Categorical存储重复字符串），减少DataFrame整体内存占用：

# 自定义schema，根据实际表结构调整
custom_schema = {
    "user_id": pl.Int32,
    "product_category": pl.Categorical,
    "price": pl.Float32
}

lf = pl.read_database(
    sql_query,
    connection=f"postgresql://{config.RDS_DB_USER}:{config.RDS_DB_PASSWORD}@{config.RDS_HOST_NAME}:{config.RDS_PORT}/{config.RDS_DB_NAME}",
    schema=custom_schema,
    engine="psycopg"
)

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