如何在Python中运行SQL查询并直接输出到CSV，避免占用内存？

我编写了一段代码运行指定目录下的.sql文件，将查询结果输出为CSV，但处理超大查询时会耗尽内存导致进程崩溃。原来的写法不管是先存DataFrame变量还是链式调用pd.read_sql(sqlFile, engine).to_csv(file_name)，本质都是先把全量结果加载到内存里。需要实现逐行直接写入文件，不预先存入全部结果。

原代码示例：

import pandas as pd
from sqlalchemy.engine import URL
from sqlalchemy import create_engine

serverName = 'FOO'
databaseName = 'BAR'

connection_string = ("Driver={SQL Server};"
            "Server=" + serverName + ";"
            "Database=" + databaseName + ";"
            "Trusted_Connection=yes;")

connection_url = URL.create(
    "mssql+pyodbc", query={"odbc_connect":connection_string })
engine = create_engine(connection_url)

sqlFile = "select 'Hello world' as 'Header'"

# 原写法：全量加载到内存再写入
pd.read_sql(sqlFile, engine).to_csv('foo.csv', header=True, index=False, sep='|')

方法1：用SQLAlchemy原生连接逐行读取写入 #

直接通过数据库连接获取游标，逐行读取查询结果并写入CSV，完全避免全量加载，内存占用极低：

import csv
from sqlalchemy.engine import URL
from sqlalchemy import create_engine

serverName = 'FOO'
databaseName = 'BAR'

connection_string = ("Driver={SQL Server};"
            "Server=" + serverName + ";"
            "Database=" + databaseName + ";"
            "Trusted_Connection=yes;")

connection_url = URL.create(
    "mssql+pyodbc", query={"odbc_connect":connection_string })
engine = create_engine(connection_url)

sql_query = "select * from large_table"  # 替换为你的超大查询

with engine.connect() as conn:
    result = conn.execute(sql_query)
    # 获取列名作为CSV表头
    columns = result.keys()
    
    with open('output.csv', 'w', newline='', encoding='utf-8') as csv_file:
        writer = csv.writer(csv_file, delimiter='|')
        # 写入表头
        writer.writerow(columns)
        # 逐行写入数据，每次仅保留一行在内存中
        for row in result:
            writer.writerow(row)

方法2：用pandas的chunksize分块读取写入 #

如果想保留pandas的操作习惯，可以用read_sql的chunksize参数分批次加载数据，控制内存占用：

import pandas as pd
from sqlalchemy.engine import URL
from sqlalchemy import create_engine

serverName = 'FOO'
databaseName = 'BAR'

connection_string = ("Driver={SQL Server};"
            "Server=" + serverName + ";"
            "Database=" + databaseName + ";"
            "Trusted_Connection=yes;")

connection_url = URL.create(
    "mssql+pyodbc", query={"odbc_connect":connection_string })
engine = create_engine(connection_url)

sql_query = "select * from large_table"
chunk_size = 10000  # 每次加载10000行，可根据内存情况调整

# 分块读取查询结果
chunks = pd.read_sql(sql_query, engine, chunksize=chunk_size)

with open('output.csv', 'w', newline='', encoding='utf-8') as csv_file:
    first_chunk = True
    for chunk in chunks:
        # 第一块写入表头，后续块跳过表头避免重复
        chunk.to_csv(csv_file, sep='|', index=False, header=first_chunk)
        first_chunk = False

为什么链式调用没用？ #

不管是pd.read_sql(...).to_csv(...)还是先存变量data = pd.read_sql(...)，read_sql都会把整个查询结果一次性加载到DataFrame中。链式调用只是省略了中间变量，内存占用完全一样，解决不了超大结果集的内存溢出问题。

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