如何用Pandas与SQLAlchemy高效可靠从PostgreSQL提取800万+行数据？

需要从一张拥有2亿+行数据的PostgreSQL数据库表中提取800万+行数据。当前实现代码如下（注：原代码存在语法小错误，已修正）：

engine = create_engine(url="MY_DB_STRING",
    echo=False,
    execution_options={'stream_results': True},
    pool_pre_ping=True,
    pool_recycle=3600
)
session = scoped_session(sessionmaker(bind=engine))
query = """
SELECT *
FROM MY_TABLE
WHERE status = True
"""
dfs = []
for chunk in pd.read_sql_query(sql=query, con=session.connection(), chunksize=500000):
    dfs.append(chunk)
combined_df = pd.concat(dfs, ignore_index=True)
session.close()

该代码在处理小型测试数据时可用，但处理真实表时耗时数小时，还可能随机卡在提取第二个数据块的环节。请问如何修改该配置，以高效且可靠地提取所有800万+行数据？

1. 移除ORM会话，直接使用引擎连接 #

SQLAlchemy的scoped_session会带来不必要的会话管理开销，批量数据提取无需依赖它。改用引擎直接连接，减少中间层损耗，同时添加超时配置避免卡死：

engine = create_engine(
    url="MY_DB_STRING",
    echo=False,
    execution_options={'stream_results': True},
    pool_pre_ping=True,
    pool_recycle=3600,
    # 设置5分钟语句超时，防止查询无响应
    connect_args={"options": "-c statement_timeout=300000"}
)

dfs = []
# 使用上下文管理器自动管理连接
with engine.connect() as conn:
    for chunk in pd.read_sql_query(sql=query, con=conn, chunksize=200000):
        dfs.append(chunk)
combined_df = pd.concat(dfs, ignore_index=True)

2. 优化查询语句与数据库索引 #

• 只查询需要的列：替换SELECT *为具体列名（如SELECT id, col_a, col_b FROM ...），减少数据传输量和内存占用，这是最直观的性能提升点。
• 给过滤字段加索引：如果status字段没有索引，执行CREATE INDEX idx_my_table_status ON MY_TABLE(status);，让WHERE条件的过滤速度大幅提升。
• 改用主键分页替代流式游标：流式读取依赖PostgreSQL服务器端游标，部分场景下会出现阻塞。改用基于主键的分页更可靠：

batch_size = 200000
last_id = 0
dfs = []

with engine.connect() as conn:
    while True:
        query = f"""
        SELECT id, col_a, col_b
        FROM MY_TABLE
        WHERE status = True AND id > {last_id}
        ORDER BY id
        LIMIT {batch_size}
        """
        chunk = pd.read_sql_query(sql=query, con=conn)
        if chunk.empty:
            break
        dfs.append(chunk)
        last_id = chunk['id'].max()

combined_df = pd.concat(dfs, ignore_index=True)

3. 调整chunksize与连接参数 #

• 减小chunksize：50万的块大小容易导致数据库端临时缓存过载，建议调整为10万-20万，平衡单次查询效率和内存压力。
• 开启网络压缩：如果数据库支持，在connect_args中添加"-c sslcompression=1"，减少网络传输的数据量：

connect_args={"options": "-c statement_timeout=300000 -c sslcompression=1"}

4. 边读边处理，避免内存累积 #

如果不需要最终合并成一个大DataFrame，可以直接将每个块写入文件，降低内存占用：

with engine.connect() as conn:
    for idx, chunk in enumerate(pd.read_sql_query(sql=query, con=conn, chunksize=100000)):
        chunk.to_csv(f"extracted_data_part_{idx}.csv", index=False)
# 后续可根据需求合并文件或分文件处理

5. 检查数据库端配置 #

• 确认PostgreSQL的work_mem设置足够处理排序和过滤操作，避免使用磁盘临时表（可临时调整：SET work_mem = '64MB';）。
• 监控数据库服务器的CPU、内存、磁盘IO负载，排除硬件瓶颈导致的卡顿。

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