如何用pandas DataFrame.to_sql()和SQLAlchemy实现异步批量插入并处理重复数据？

在FastAPI路由中使用pandas的DataFrame.to_sql()结合SQLAlchemy异步会话实现基因数据批量插入，需处理表中已存在的重复数据。当前实现代码及遇到的问题如下：

批量插入函数 #

async def bulk_insert_genes(data_sets: pd.DataFrame):
    """Bulk insert genes from bucket to postgresDB.

    Args:
        data_sets: Dataframe of paths to the extraction datasets.
    """
    async with session_maker() as session:
        for _, row in data_sets.iterrows():
            data = generate_gene_df(row)
            conn = await session.connection()
            await conn.run_sync(
                lambda sync_conn, data=data: data.to_sql(
                    name=GenePos.__tablename__,
                    con=sync_conn,
                    if_exists="append",
                    index=False,
                    # method="multi",
                    chunksize=CHUNK_SIZE,
                    method=insert_on_duplicate,
                ),
            )
            await session.commit()

重复数据处理函数及报错 #

参考实现的自定义插入方法运行时抛出异步不支持错误：AttributeError: 'PGCompiler_asyncpg' object has no attribute 'visit_on_duplicate_key_update'

def insert_on_duplicate(table, conn, keys, data_iter):  # noqa: ANN001
    """Insert data into table with on duplicate key update."""
    insert_stmt = insert(table.table).values(list(data_iter))
    on_duplicate_key_stmt = insert_stmt.on_duplicate_key_update(insert_stmt.inserted)
    conn.execute(on_duplicate_key_stmt)

注：generate_gene_df()用于生成适配数据库表格式的DataFrame，需更优的实现方式及重复数据处理方案。

1. 错误根源 #

on_duplicate_key_update是MySQL专属语法，PostgreSQL对应的是ON CONFLICT ... DO UPDATE逻辑，且异步SQLAlchemy的PostgreSQL编译器未实现visit_on_duplicate_key_update方法，因此触发报错。

2. 适配PostgreSQL的重复数据处理函数 #

修改自定义插入方法，使用PostgreSQL原生的ON CONFLICT语法，需提前明确表的主键/唯一约束列（示例假设唯一约束为gene_id，请根据实际表结构修改）：

from sqlalchemy import insert

def insert_on_conflict_do_update(table, conn, keys, data_iter):
    """PostgreSQL批量插入时处理重复数据：存在则更新，不存在则插入"""
    data_list = list(data_iter)
    insert_stmt = insert(table.table).values(data_list)
    
    # 替换为你的表的主键/唯一约束列
    conflict_columns = ['gene_id']
    # 构建需要更新的字段（排除主键/唯一约束列）
    update_dict = {
        key: insert_stmt.inserted[key]
        for key in keys
        if key not in conflict_columns
    }
    
    on_conflict_stmt = insert_stmt.on_conflict_do_update(
        index_elements=conflict_columns,
        set_=update_dict
    )
    conn.execute(on_conflict_stmt)

3. 优化批量插入逻辑 #

原代码逐行循环处理并频繁提交事务，可优化为合并数据后批量处理，减少数据库交互开销：

async def bulk_insert_genes(data_sets: pd.DataFrame):
    """Bulk insert genes from bucket to postgresDB."""
    async with session_maker() as session:
        # 合并所有待插入的基因数据
        gene_data_frames = []
        for _, row in data_sets.iterrows():
            gene_df = generate_gene_df(row)
            gene_data_frames.append(gene_df)
        combined_df = pd.concat(gene_data_frames, ignore_index=True)
        
        conn = await session.connection()
        await conn.run_sync(
            lambda sync_conn: combined_df.to_sql(
                name=GenePos.__tablename__,
                con=sync_conn,
                if_exists="append",
                index=False,
                chunksize=CHUNK_SIZE,
                method=insert_on_conflict_do_update,
            ),
        )
        await session.commit()

4. 额外性能优化建议 #

• 若generate_gene_df支持批量处理，直接传入整个data_sets生成合并后的DataFrame，避免循环调用
• 根据数据库性能调整chunksize参数（建议范围1000-5000条/块）
• 确保表的主键/唯一约束已正确创建，否则ON CONFLICT逻辑不会生效

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