如何用SQLAlchemy将大文件（50-100GB）导入关联数据库？

看起来你遇到了SQLAlchemy循环提交导致的性能瓶颈，同时还要处理关联表的唯一约束和ID关联问题。我来帮你梳理下最佳实践，把你的导入速度提上去，同时解决序列重复和关联的问题。

核心问题分析 #

你的原代码每次循环都执行session.commit()，这会导致大量事务开销——SQLite每次提交都要刷新磁盘日志，频繁提交是性能差的主要原因。另外，处理重复序列的try/except commit方式不仅低效，还可能导致逻辑漏洞。

优化方案步骤 #

1. 开启SQLite WAL模式提升写入性能

SQLite默认的日志模式是DELETE，开启WAL（Write-Ahead Logging）可以大幅提升写入速度，尤其是批量操作时。在创建引擎后执行PRAGMA journal_mode=WAL。

2. 批量处理，减少事务次数

把所有数据收集后，一次性处理序列的去重、插入，再批量插入关联的Protein记录，最后只做一次commit。

3. 复用已存在的序列ID

因为Protein_sequence.prot_seq是唯一约束，先查询数据库中已存在的序列，避免重复插入，同时建立序列到ID的映射，方便后续关联Protein表。

4. 修复FASTA解析的隐藏bug

原parse_fasta函数中，prot字典是在循环外创建的，yield的是同一个字典的引用，会导致所有记录被最后一条覆盖。要把字典创建移到循环内部。

优化后的完整代码 #

import collections
import re
import Bio.SeqIO
import sqlalchemy
from sqlalchemy import ForeignKey, UniqueConstraint
from sqlalchemy import Column, Float, Integer, String, Text, DateTime
from sqlalchemy.sql import func
from sqlalchemy.orm import relationship
from sqlalchemy.ext.declarative import declarative_base

Base = declarative_base()

class Protein_sequence(Base):
    __tablename__ = 'protein_sequence'
    prot_seq_id = Column(Integer, primary_key=True)
    prot_seq = Column(Text, unique=True)
    protein_annotation = relationship('Protein', back_populates='protein_sequence')

class Protein(Base):
    __tablename__ = 'protein_annotation'
    prot_id = Column(Integer, primary_key=True)
    prot_seq_id = Column(Integer, ForeignKey('protein_sequence.prot_seq_id'))
    prot_acc = Column(Text, unique=True)
    prot_name = Column(Text)
    prot_db = Column(Text)  # 补充原代码中缺失的字段，匹配parse_fasta的输出
    prot_gi = Column(Text)  # 补充原代码中缺失的字段
    protein_sequence = relationship('Protein_sequence', back_populates='protein_annotation')

def parse_fasta(path, prot_db='unknown', taxon_name=None, taxon_id=None):
    """Parsing a fasta file (UniProt or NCBInr)."""
    for record in Bio.SeqIO.parse(path, 'fasta'):
        # 每次循环新建字典，避免引用覆盖问题
        prot = collections.OrderedDict()
        prot['seq'] = str(record.seq)
        desc = record.description
        
        gi_num = re.findall('^gi\|([0-9]+)(?:\||\s|$)', desc)
        if gi_num:
            prot['prot_gi'] = gi_num[0]
            desc = re.sub('^gi\|([0-9]+)(?:\||\s|$)', '', desc)
        
        # 修复prot_db的处理逻辑，保留传入的默认值
        matched_db = re.findall('^([^|]+)\|', desc)
        if matched_db:
            prot['prot_db'] = matched_db[0]
        else:
            prot['prot_db'] = prot_db
        
        prot_acc = re.findall('^[^|]+\|([^ ]+)', desc)[0]
        prot['prot_acc'] = prot_acc
        
        prot_name_match = re.findall('^[^ ]+ (.+)', desc)
        if prot_name_match:
            prot['prot_name'] = prot_name_match[0]
        else:
            prot['prot_name'] = ''  # 处理无名称的异常情况
        
        yield prot

def prot_db_from_fasta():
    """Create tables in SQLite database. Input fasta file."""
    db = 'sqlite:///proteomic.db'
    # 创建引擎，配置SQLite优化参数
    engine = sqlalchemy.create_engine(
        db,
        connect_args={"check_same_thread": False},  # 支持多线程访问
        pool_pre_ping=True
    )
    # 开启WAL模式提升写入性能
    with engine.connect() as conn:
        conn.execute(sqlalchemy.text("PRAGMA journal_mode=WAL;"))
    
    Base.metadata.create_all(engine)
    Session = sqlalchemy.orm.sessionmaker(bind=engine)
    session = Session()
    
    try:
        p = 'prot.fasta'
        # 先收集所有蛋白数据，方便批量处理
        all_prots = list(parse_fasta(p))
        if not all_prots:
            print("No proteins found in the FASTA file.")
            return
        
        # 步骤1：提取所有唯一序列，去重
        unique_seqs = list({prot['seq'] for prot in all_prots})
        
        # 步骤2：查询数据库中已存在的序列，建立序列到ID的映射
        existing_seq_query = session.query(Protein_sequence.prot_seq, Protein_sequence.prot_seq_id)\
                                    .filter(Protein_sequence.prot_seq.in_(unique_seqs))
        seq_to_id = {seq: seq_id for seq, seq_id in existing_seq_query.all()}
        
        # 步骤3：批量插入新的序列
        new_seq_objects = [Protein_sequence(prot_seq=seq) for seq in unique_seqs if seq not in seq_to_id]
        if new_seq_objects:
            session.add_all(new_seq_objects)
            session.flush()  # 刷新会话，获取新插入的prot_seq_id
            # 更新序列ID映射
            for seq_obj in new_seq_objects:
                seq_to_id[seq_obj.prot_seq] = seq_obj.prot_seq_id
        
        # 步骤4：批量创建Protein记录
        protein_objects = []
        # 获取Protein表的有效字段名
        protein_columns = {c.name for c in Protein.__table__.columns}
        for prot in all_prots:
            # 过滤出Protein表存在的字段
            annotation_data = {key: value for key, value in prot.items() if key in protein_columns}
            # 关联对应的序列ID
            annotation_data['prot_seq_id'] = seq_to_id[prot['seq']]
            protein_objects.append(Protein(**annotation_data))
        
        if protein_objects:
            session.add_all(protein_objects)
        
        # 最后一次性提交所有更改
        session.commit()
        print(f"Successfully imported {len(all_prots)} proteins, {len(new_seq_objects)} new sequences.")
    
    except Exception as e:
        session.rollback()
        print(f"Error during import: {str(e)}")
        raise
    finally:
        session.close()

# run function to insert data into database
prot_db_from_fasta()

关键优化点说明 #

• WAL模式：开启后SQLite支持并发读写，写入性能提升数倍，接近原生sqlite3的速度。
• 批量操作：所有插入操作都在一个事务中完成，只做一次commit，避免了频繁事务的磁盘开销。
• 序列去重与复用：先查询已存在的序列，只插入新序列，同时建立ID映射，确保Protein表能正确关联到对应的序列ID。
• 修复解析bug：把prot字典移到循环内部，避免了所有记录被最后一条覆盖的问题。

额外性能建议 #

• 如果FASTA文件极大（比如几个G），可以分批次处理（比如每10000条记录处理一次），避免内存占用过高。
• 对于SQLite，还可以设置PRAGMA synchronous=NORMAL进一步提升速度（但可能降低数据安全性，适合导入时临时使用）。

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