Python Concurrent Futures迁移超大PostgreSQL表至SQL Server内存溢出问题

我正在开发Python代码，将50GB以上的PostgreSQL表动态迁移至SQL Server，采用SQLAlchemy流式查询结合ThreadPoolExecutor批量插入，每个线程创建独立连接（同一连接多插入会导致SQL Server繁忙）。

5-10GB的小表运行正常，但超大表会出现内存持续攀升直至耗尽的情况，内存使用曲线持续上升。已尝试清理存储数据的列表、显式删除线程Future，但仍无法解决内存溢出问题，核心代码片段如下：

主函数代码 #

for table in tables:
    table_name = table[0]
    print("table :", table_name)
    # Connect to PostgreSQL
    pg_conn = psycopg2.connect(host=pg_server, port=pg_port, database=pg_db, user=pg_user, password=pg_password)
    pg_cursor = pg_conn.cursor()
    # Get column names and datatypes from PostgreSQL
    pg_cursor.execute(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_schema = %s AND table_name = %s", (schema_name, table_name))
    pg_columns = pg_cursor.fetchall()
    pg_cursor.close()
    pg_conn.close()

    # Connect to SQL Server
    sql_conn = pyodbc.connect('DRIVER={SQL Server};SERVER=' + sql_server + ';DATABASE=' + sql_db + ';UID=' + sql_user + ';PWD=' + sql_password)
    sql_cursor = sql_conn.cursor()

    sq_columns = [(column[0], dtmapping[column[1]]) for column in pg_columns]
    sq_column_definitions = ', '.join([f"[{column[0]}] {column[1]}" for column in sq_columns])
    sq_cols = [column[0] for column in sq_columns]
    def casting(column):
        if column[1] in ['json','bjson']:
            return f'"{column[0]}"::text as "{column[0]}"'
        else:
            return f'"{column[0]}"'
    column_names = [casting(column) for column in pg_columns] 
    column_names_str = ', '.join(column_names)
    if schema_name == 'public':
        schema_name = 'dbo'
    # Create schema and table in SQL Server
    sql_cursor.execute(f"IF NOT EXISTS (SELECT * FROM sys.schemas WHERE name='{schema_name}') BEGIN EXEC('CREATE SCHEMA [{schema_name}]') END")
    sql_cursor.execute(f"IF NOT EXISTS (Select * from information_schema.tables where TABLE_NAME ='{table_name}' and TABLE_SCHEMA = '{schema_name}' AND TABLE_type = 'BASE TABLE') BEGIN CREATE TABLE [{schema_name}].[{table_name}] ({sq_column_definitions}) END")
    sql_cursor.execute(f"TRUNCATE TABLE {schema_name}.{table_name}")

    # Commit the changes
    sql_conn.commit()
    sql_cursor.close()
    sql_conn.close()

    def yieldmany(exe,chunk_size,chunk_cnt):
        while 'batch not empty':
            batch = exe.fetchmany(chunk_size)
            if not batch:
                break
            chunk_cnt += 1
            yield pd.DataFrame(batch) , chunk_cnt
    
    #Connect to postgres using SQLAlchemy Source
    enginer = create_engine(f'postgresql+psycopg2://{pg_user}:{pg_password}@{pg_server}/{pg_db}',echo=False, echo_pool=False, pool_pre_ping=True)
    enginer.execution_options(stream_results=True)
    sqlar_conn = enginer.connect()
    
    print(f"SELECT {column_names_str} FROM {schema_name}.{table_name} limit 1000")
    query = text(f"SELECT {column_names_str} FROM {schema_name}.{table_name}")
    total_chunks = 0
    chunk_cnt = 0
    with ThreadPoolExecutor(max_workers=20) as executor:
        futuress = [executor.submit(worker, chunk, table_name, schema_name,sq_cols,chunk_cnt) for chunk, chunk_cnt in yieldmany(enginer.connect().execution_options(stream_results=True).execute(query),chunk_size,chunk_cnt)]
        for future in futures.as_completed(futuress):
            result, chunk_cnt = future.result()
            del futuress[future]
            total_chunks += chunk_cnt
            print(result)
            
    print(f'{table_name} completed')
    completed_tables.add(table_name)
    logger.info(f'completed tables : {completed_tables}')
    if schema_name == 'dbo':
        schema_name = 'public'
    chunk_cnt = 0
    sqlar_conn.close()
    # exe.close()
    enginer.dispose()

Worker函数代码 #

def worker(chunk: pd.DataFrame, table_name: str, schema_name: str, sq_cols: list, chunk_cnt: int):
    logger.info(f'chunking {chunk_cnt} part for {table_name} started')
    engine_url = f'mssql+pyodbc://{sql_user}:{sql_password}@{sql_server}/{sql_db}?trusted_connection=yes&driver=ODBC+Driver+17+for+SQL+Server'
    with create_engine(engine_url).connect() as sqlaw_conn:
        sqlaw_conn.execution_options(autocommit=True)
        try:
            chunk.to_sql(table_name, sqlaw_conn, schema=schema_name, if_exists='append', index=False)
            logger.info(f"{table_name} {chunk_cnt} chunked")
        except Exception as e:
            print(f"Error inserting data: {e}")
            problem_tables.add(table_name)
            logger.info(f"Error inserting data for {table_name}: {e}")
            logger.error(f'problem tables : {problem_tables}')
            print("problem_tables : {}".format(problem_tables))
    logger.info(f'problem tables : {problem_tables}')
    return f"{table_name} {chunk_cnt} chunked", chunk_cnt

问题根源与优化方案 #

核心问题分析 #

1. 流式查询完全失效：用列表推导式一次性遍历完yieldmany生成器，会把超大表的所有chunk数据提前加载到内存，完全违背流式查询的设计初衷。
2. 线程池任务提交逻辑错误：一次性创建所有Future并存入列表，导致内存瞬间被所有chunk数据占满。
3. 重复创建SQLAlchemy引擎：Worker函数每次执行都新建引擎，带来额外内存开销，且无法复用连接池。
4. DataFrame内存未优化：默认方式创建的DataFrame对大字段（如JSON转文本）的内存占用未做控制。

针对性优化 #

1. 修正流式任务提交逻辑

不要一次性提交所有任务，控制并发数，避免全量加载数据：

# 替换原线程池代码块
import concurrent.futures

with ThreadPoolExecutor(max_workers=20) as executor:
    futures = set()
    total_chunks = 0
    chunk_cnt = 0
    # 逐个生成chunk，控制并发数
    for chunk, chunk_cnt in yieldmany(enginer.connect().execution_options(stream_results=True).execute(query), chunk_size, chunk_cnt):
        # 并发数达到上限时，等待部分任务完成再提交新任务
        while len(futures) >= 20:
            done, futures = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_COMPLETED)
            for future in done:
                result, cnt = future.result()
                total_chunks += cnt
                print(result)
        # 提交新任务
        future = executor.submit(worker, chunk, table_name, schema_name, sq_cols, chunk_cnt)
        futures.add(future)
    # 等待剩余任务完成
    for future in concurrent.futures.as_completed(futures):
        result, cnt = future.result()
        total_chunks += cnt
        print(result)

2. 复用SQL Server引擎

提前创建一次引擎，避免重复初始化：

# 在主函数开头或全局范围创建SQL Server引擎
sql_server_engine = create_engine(
    f'mssql+pyodbc://{sql_user}:{sql_password}@{sql_server}/{sql_db}?trusted_connection=yes&driver=ODBC+Driver+17+for+SQL+Server',
    pool_size=20,  # 匹配线程池大小
    max_overflow=0
)

# 修改Worker函数
def worker(chunk: pd.DataFrame, table_name: str, schema_name: str, sq_cols: list, chunk_cnt: int):
    logger.info(f'chunking {chunk_cnt} part for {table_name} started')
    with sql_server_engine.connect() as sqlaw_conn:
        sqlaw_conn.execution_options(autocommit=True)
        try:
            chunk.to_sql(table_name, sqlaw_conn, schema=schema_name, if_exists='append', index=False)
            logger.info(f"{table_name} {chunk_cnt} chunked")
        except Exception as e:
            print(f"Error inserting data: {e}")
            problem_tables.add(table_name)
            logger.info(f"Error inserting data for {table_name}: {e}")
            logger.error(f'problem tables : {problem_tables}')
            print("problem_tables : {}".format(problem_tables))
    logger.info(f'problem tables : {problem_tables}')
    return f"{table_name} {chunk_cnt} chunked", chunk_cnt

3. 优化DataFrame内存占用

指定字段数据类型，减少内存消耗：

def yieldmany(exe, chunk_size, chunk_cnt):
    # 构建PostgreSQL类型到pandas类型的映射
    dtype_mapping = {}
    for col in pg_columns:
        pg_type = col[1]
        if pg_type == 'integer':
            dtype_mapping[col[0]] = 'int32'
        elif pg_type == 'bigint':
            dtype_mapping[col[0]] = 'int64'
        elif pg_type in ['numeric', 'decimal']:
            dtype_mapping[col[0]] = 'float64'
        elif pg_type in ['json', 'bjson']:
            dtype_mapping[col[0]] = 'string'
        else:
            dtype_mapping[col[0]] = 'object'
    
    while True:
        batch = exe.fetchmany(chunk_size)
        if not batch:
            break
        chunk_cnt += 1
        # 指定dtype创建DataFrame
        df = pd.DataFrame(batch, columns=[col[0] for col in pg_columns], dtype=dtype_mapping)
        yield df, chunk_cnt

4. 其他细节优化

• 移除主函数中未使用的sqlar_conn变量，减少无效内存占用；
• 导入gc模块，在每次处理完任务后调用gc.collect()，主动触发垃圾回收；
• 调整日志级别，避免冗余的大内容日志输出占用内存。

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