如何加速超6亿条数据的SQLite rsID查询？

问题背景 #

本地dbSNP数据库存储了约6亿条条目，仅保留rsID与基因组位置相关信息（rsid、chrom、pos、ref、alt）。此前用Pandas直接查询速度过慢，遂将数据导入SQLite，但当前查询无论条目数量多少，均耗时约1.5分钟，急需优化查询性能。

数据导入代码 #

当初用于将VCF.gz文件导入SQLite的代码：

import sqlite3
import gzip
import csv

rsid_db = sqlite3.connect('rsid.db')
rsid_cursor = rsid_db.cursor()

rsid_cursor.execute(
"""
CREATE TABLE rsids (
rsid TEXT,
chrom TEXT,
pos INTEGER,
ref TEXT,
alt TEXT
)
"""
)

with gzip.open('00-All.vcf.gz', 'rt') as vcf:
    reader = csv.reader(vcf, delimiter="\t")
    i = 0
    for row in reader:
        if not ''.join(row).startswith('#'):
            rsid_cursor.execute(
            f"""
            INSERT INTO rsids (rsid, chrom, pos, ref, alt)
            VALUES ('{row[2]}', '{row[0]}', '{row[1]}', '{row[3]}', '{row[4]}');
            """
            )
            i += 1
            if i % 1000000 == 0:
                print(f'{i} entries written')
                rsid_db.commit()

rsid_db.commit()
rsid_db.close()

当前查询代码 #

目前用于批量查询rsID的代码：

import sqlite3
import pandas as pd

def query_rsid(rsid_list,
               rsid_db_path='rsid.db'):

    with sqlite3.connect(rsid_db_path) as rsid_db:
        rsid_cursor = rsid_db.cursor()
        rsid_cursor.execute(
        f"""
        SELECT * FROM rsids
        WHERE rsid IN ('{"', '".join(rsid_list)}');
        """
        )
        query = rsid_cursor.fetchall()

    return query

优化方案 #

1. 为rsid字段添加索引（核心优化） #

当前查询速度慢的根本原因是无索引导致全表扫描——6亿条数据的全表扫描必然耗时极长。解决方法是给rsid字段创建唯一索引（每个rsID都是唯一的）：

方式一：事后添加索引

打开SQLite连接，执行以下SQL命令：

CREATE UNIQUE INDEX idx_rsid ON rsids(rsid);

注：创建索引需要一定时间（取决于硬件），但完成后后续所有基于rsid的查询速度会提升几个数量级。

方式二：重建表时设为主键

如果允许重新导入数据，可以直接将rsid设为主键，SQLite会自动创建唯一索引：

CREATE TABLE rsids (
rsid TEXT PRIMARY KEY,
chrom TEXT,
pos INTEGER,
ref TEXT,
alt TEXT
);

2. 优化查询代码（参数化+高效读取） #

原查询代码用字符串拼接生成IN条件，存在SQL注入风险，且不利于SQLite优化查询计划。改用参数化查询，同时用Pandas的read_sql_query直接读取结果，效率更高：

import sqlite3
import pandas as pd

def query_rsid(rsid_list, rsid_db_path='rsid.db'):
    # 生成与rsid数量匹配的占位符
    placeholders = ', '.join(['?'] * len(rsid_list))
    query_sql = f"SELECT rsid, chrom, pos, ref, alt FROM rsids WHERE rsid IN ({placeholders})"
    
    with sqlite3.connect(rsid_db_path) as conn:
        # 启用WAL模式提升并发性能
        conn.execute('PRAGMA journal_mode = WAL')
        # 用Pandas读取结果，比fetchall更高效
        result_df = pd.read_sql_query(query_sql, conn, params=rsid_list)
    
    # 按需返回格式：可返回字典列表或直接返回DataFrame
    return result_df.to_dict('records')

3. 调整SQLite配置提升性能 #

在连接数据库时设置以下参数，利用更多内存缓存并优化读写模式：

conn = sqlite3.connect('rsid.db')
# 设置2GB内存缓存（负号表示单位为KB）
conn.execute('PRAGMA cache_size = -2097152')
# 启用写前日志模式，提升读写并发
conn.execute('PRAGMA journal_mode = WAL')
# 降低同步级别，平衡性能与安全性
conn.execute('PRAGMA synchronous = NORMAL')
# 临时表存入内存，加快查询
conn.execute('PRAGMA temp_store = MEMORY')

4. 补充：优化数据导入效率（若需重新导入） #

原导入代码单条插入效率极低，6亿条数据会花费大量时间。改用executemany批量插入，速度能提升10-100倍：

import sqlite3
import gzip
import csv

rsid_db = sqlite3.connect('rsid.db')
rsid_db.execute('PRAGMA journal_mode = WAL')
rsid_cursor = rsid_db.cursor()

rsid_cursor.execute(
"""
CREATE TABLE rsids (
rsid TEXT PRIMARY KEY,
chrom TEXT,
pos INTEGER,
ref TEXT,
alt TEXT
)
"""
)

batch_size = 10000
batch_data = []

with gzip.open('00-All.vcf.gz', 'rt') as vcf:
    reader = csv.reader(vcf, delimiter="\t")
    i = 0
    for row in reader:
        if not ''.join(row).startswith('#'):
            batch_data.append((row[2], row[0], int(row[1]), row[3], row[4]))
            i += 1
            if i % batch_size == 0:
                rsid_cursor.executemany(
                    "INSERT INTO rsids (rsid, chrom, pos, ref, alt) VALUES (?, ?, ?, ?, ?)",
                    batch_data
                )
                rsid_db.commit()
                batch_data = []
                print(f'{i} entries written')
    # 插入剩余数据
    if batch_data:
        rsid_cursor.executemany(
            "INSERT INTO rsids (rsid, chrom, pos, ref, alt) VALUES (?, ?, ?, ?, ?)",
            batch_data
        )
        rsid_db.commit()

rsid_db.close()

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