从大文件中检索行的更有效方法

Question

我有一个 ID 为 1,786,916 条记录的数据文件，我想从另一个包含大约 480 万条记录（在本例中为 DNA 序列，但基本上只是纯文本）的文件中检索相应的记录。我编写了一个 python 脚本来执行此操作，但它需要很长时间才能运行（第 3 天，它只完成了 12%）。由于我是 python 的相对新手，我想知道是否有人建议让这更快。

以下是带有 ID 的数据文件示例（示例中的 ID 为 ANCH889-10）：

ANICH889-10 k__Animalia; p__Arthropoda; c__Insecta; o__Lepidoptera; f__Psychidae; g__Ardiosteres; s__Ardiosteres sp. ANIC9
ARONW984-15 k__Animalia; p__Arthropoda; c__Arachnida; o__Araneae; f__Clubionidae; g__Clubiona; s__Clubiona abboti

这是包含记录的第二个文件的示例：

>ASHYE2081-10|Creagrura nigripesDHJ01|COI-5P|HM420985
ATTTTATACTTTTTATTAGGAATATGATCAGGAATAATTGGTCTTTCAATAAGAATCATTATCCGTATTGAATTAAGAAATCCAGGATCTATTATTAATAATGACCAAATTTATAATTCATTAATTACTATACACGCACTATTAATAATTTTTTTTTTAGTTATACCTGTAATAATTGGAGGATTTGGAAATTGATTAATTCCTATTATAATTGGAGCCCCAGATATAGCATTTCCACGAATAAACAATCTTAGATTTTGATTATTAATCCCATCAATTTTCATATTAATATTAAGATCAATTACTAATCAAGGTGTAGGAACAGGATGAACAATATATCCCCCATTATCATTAAATATAAATCAAGAAGGAATATCAATAGATATATCAATTTTTTCTTTACATTTAGCAGGAATATCCTCAATTTTAGGATCAATTAATTTCATTTCAACTATTTTAAATATAAAATTTATTAATTCTAATTATGATCAATTAACTTTATTTTCATGATCAATTCTAATTACTACTATTTTATTATTACTAGCAGTCCCTGTATTAGCAGGAGCAATTACTATAATTTTAACTGATCGAAATTTAAATACTTCTTTTTTTGATCCTAGAGGAGGAGGAGATCCAATTT-----------------
>BCISA145-10|Hemiptera|COI-5P
AACTCTATACTTTTTACTAGGATCCTGGGCAGGAATAGTAGGAACATCATTAAGATGAATAATCCGAATTGAACTAGGACAACCTGGATCTTTTATTGGAGATGACCAAACTTATAATGTAATTGTAACTGCCCACGCATTTGTAATAATTTTCTTTATAGTTATACCAATTATAATTGGAGGATTTGGAAATTGATTAATTCCCTTAATAATTGGAGCACCCGATATAGCATTCCCACGAATGAATAACATAAGATTTTGATTGCTACCACCGTCCCTAACACTTCTAATCATAAGTAGAATTACAGAAAGAGGAGCAGGAACAGGATGAACAGTATACCCTCCATTATCCAGAAACATCGCCCATAGAGGAGCATCTGTAGATTTAGCAATCTTTTCCCTACATCTAGCAGGAGTATCATCAATTTTAGGAGCAGTTAACTTCATTTCAACAATTATTAATATACGACCAGCAGGAATAACCCCAGAACGAATCCCATTATTTGTATGATCTGTAGGAATTACAGCACTACTACTCCTACTTTCATTACCCGTACTAGCAGGAGCCATTACCATACTCTTAACTGACCGAAACTTCAATACTTCTTTTTTTGACCCTGCTGGAGGAGGAGATCCCATCCTATATCAACATCTATTC

但是在第二个文件中，DNA 序列被分成几行，而不是单行，而且它们的长度并不总是相同。

编辑

这是我想要的输出：

>ANICH889-10
GGGATTTGGTAATTGATTAGTTCCTTTAATA---TTGGGGGCCCCTGACATAGCTTTTCCTCGTATAAATAATATAAGATTTTGATTATTACCTCCCTCTCTTACATTATTAATTTCAAGAAGAATTGTAGAAAATGGAGCTGGGACTGGATGAACTGTTTACCCTCCTTTATCTTCTAATATCGCCCATAGAGGAAGCTCTGTAGATTTA---GCAATTTTCTCTTTACATTTAGCAGGAATTTCTTCTATTTTAGGAGCAATTAATTTTATTACAACAATTATTAATATACGTTTAAATAATTTATCTTTCGATCAAATACCTTTATTTGTTTGAGCAGTAGGAATTACAGCATTTTTACTATTACTTTCTTTACCTGTATTAGCTGGA---GCTATTACTATATTATTAACT---------------------------------------------------------------------------
>ARONW984-15
TGGTAACTGATTAGTTCCATTAATACTAGGAGCCCCTGATATAGCCTTCCCCCGAATAAATAATATAAGATTTTGACTTTTACCTCCTTCTCTAATTCTTCTTTTATCAAGGTCTATTATNGAAAATGGAGCA---------GGAACTGGCTGAACAGTTTACCCTCCCCTTTCTTNTAATATTTCCCATGCTGGAGCTTCTGTAGATCTTGCAATCTTTTCCCTACACCTAGCAGGTATTTCCTCAATCCTAGGGGCAGTTAAT------TTTATCACAACCGTAATTAACATACGCTCTAGAGGAATTACATTTGATCGAATGCCTTTATTTGTATGATCTGTATTAATTACAGCTATTCTTCTACTACTCTCCCTCCCAGTATTAGCAGGGGCTATTACAATACTACTCACAGACCGAAATTTAAAT-----------------------------------

这是我为此编写的 python 脚本：

from Bio import SeqIO
from Bio.Seq import Seq
import csv
import sys

#Name of the datafile
Taxonomyfile = "02_Arthropoda_specimen_data_less.txt"

#Name of the original sequence file
OrigTaxonSeqsfile = "00_Arthropoda_specimen.fasta"

#Name of the output sequence file
f4 = open("02_Arthropoda_specimen_less.fasta", 'w')

#Reading the datafile and extracting record IDs   
TaxaKeep = []
with open(Taxonomyfile, 'r') as f1:
    datareader = csv.reader(f1, delimiter='\t')
    for item in datareader:
        TaxaKeep.append(item[0])
    print(len(TaxaKeep))    

#Filtering sequence file to keep only those sequences with the desired IDs
datareader = SeqIO.parse(OrigTaxonSeqsfile, "fasta")
for seq in datareader:
    for item in TaxaKeep:
        if item in seq.id:
            f4.write('>' + str(item) + '\n')
            f4.write(str(seq.seq) + '\n')

我认为这里的问题在于，我正在为 480 万条记录中的每一条循环遍历 170 万条记录名称的列表。我想过为这 480 万条记录制作一本字典之类的东西，但我不知道该怎么做。有什么建议（包括非 python 建议）？

谢谢！

Answer 1

您的推理是正确的，即使用两个嵌套的 for 循环，您将花费时间为 4.8 million * 1.7 million 重复执行单个操作。

这就是为什么我们将使用 SQLite 数据库来存储 OrigTaxonSeqsfile 中包含的所有信息。为什么选择 SQLite？因为

• SQLite 内置于 Python
• SQLite 支持索引

我无法开始解释 CS 理论，但是在像您这样的案例中搜索数据时，索引是上帝派来的。

一旦数据被索引，您只需在数据库中查找来自Taxonomyfile 的每条记录ID，并将其写入您的f4 最终输出文件。

以下代码应该可以按照您的意愿工作，它具有以下优点：

• 显示您在处理的行数方面取得的进展
• 只需要 Python 3，不需要 Bio 库
• 使用生成器，因此不必一次将文件全部读入内存
• 不依赖于 list/set/dict，因为（在这种特殊情况下）它们可能会消耗过多的 RAM

这是代码

import sqlite3
from itertools import groupby
from contextlib import contextmanager

Taxonomyfile = "02_Arthropoda_specimen_data_less.txt"
OrigTaxonSeqsfile = "00_Arthropoda_specimen.fasta"

@contextmanager
def create_db(file_name):
    """ create SQLite db, works as context manager so file is closed safely"""
    conn = sqlite.connect(file_name, isolation_level="IMMEDIATE")
    cur = conn.connect()
    cur.execute("""
        CREATE TABLE taxonomy
        ( _id INTEGER PRIMARY KEY AUTOINCREMENT
        , record_id TEXT NOT NULL
        , record_extras TEXT
        , dna_sequence TEXT
        );
        CREATE INDEX idx_taxn_recID ON taxonomy (record_id);
    """)
    yield cur
    conn.commit()
    conn.close()
    return

def parse_fasta(file_like):
    """ generate that yields tuple containing record id, extra info
    in tail of header and the DNA sequence with newline characters
    """
    # inspiration = https://www.biostars.org/p/710/
    try:
        from Bio import SeqIO
    except ImportError:
        fa_iter = (x[1] for x in groupby(file_like, lambda line: line[0] == ">"))
        for header in fa_iter:
            # remove the >
            info = header.__next__()[1:].strip()
            # seprate record id from rest of the seqn info
            x = info.split('|')
            recID, recExtras = x[0], x[1:]
            # build the DNA seq using generator
            sequence = "".join(s.strip() for s in fa_iter.__next__())
            yield recID, recExtras, sequence
    else:
        fasta_sequences = SeqIO.parse(file_like, 'fasta')
        for fasta in fasta_sequences:
            info, sequence = fasta.id, fasta.seq.tostring()
            # seprate record id from rest of the seqn info
            x = info.split('|')
            recID, recExtras = x[0], x[1:]
            yield recID, recExtras, sequence
    return

def prepare_data(txt_file, db_file):
    """ put data from txt_file into db_file building index on record id """
    i = 0
    src_gen = open(txt_file, mode='rt')
    fasta_gen = parse_fasta(src_gen)
    with create_db(db_file) as db:
        for recID, recExtras, dna_seq in fasta_gen:
            db.execute("""
                INSERT INTO taxonomy
                (record_id, record_extras, dna_sequence) VALUES (?,?,?)
                """,
                [recID, recExtras, dna_seq]
            )
            if i % 100 == 0:
                print(i, 'lines digested into sql database')
    src_gen.close()
    return

def get_DNA_seq_of(recordID, src):
    """ search for recordID in src database and return a formatted string """
    ans = ""
    exn = src.execute("SELECT * FROM taxonomy WHERE record_id=?", [recordID])
    for match in exn.fetchall():
        a, b, c, dna_seq = match
        ans += ">%s\n%s\n" % (recordID, dna_seq)
    return ans

def main():
    # first of all prepare an optimized database
    db_file = txt_file + ".sqlite"
    prepare_data(OrigTaxonSeqsfile)
    # now start searching and writing
    progress = 0
    db = sqlite3.connect(db_file)
    cur = db.cursor()
    out_file = open("02_Arthropoda_specimen_less.fasta", 'wt')
    taxa_file = open(Taxonomyfile, 'rt')
    with taxa_file, out_file:
        for line in taxa_file:
            question = line.split("\t")[0]
            answer = get_DNA_seq_of(question, cur)
            out_file.write(answer)
            if progress % 100 == 0:
                print(progress, 'lines processed')
    db.close()

if __name__ == '__main__':
    main()

如有任何疑问，请随时询问。
如果您遇到任何错误或输出不符合预期，请给我发送 200 行示例，分别为 Taxonomyfile 和 OrigTaxonSeqsfile，我会更新代码。

---

速度增益

以下是一个粗略的估计，只讨论磁盘 I/O，因为这是最慢的部分。

让a = 4.8 million 和b = 1.7 million。

在旧方法中，您必须执行磁盘 I/O a * b 即 81600 亿 次。

在我的方法中，一旦您进行索引（即 2*a 次），您必须搜索 170 万条记录。所以在我的方法中，总时间是2 * (a + b)，即 1300 万磁盘 I/O，这也不小，但这种方法比 快 60 万倍

为什么不dict()？

如果发现我使用过多的 CPU/RAM，我会被老板和教授责骂。如果您拥有该系统，则更简单的基于 dict 的方法是：

from itertools import groupby

Taxonomyfile = "02_Arthropoda_specimen_data_less.txt"
OrigTaxonSeqsfile = "00_Arthropoda_specimen.fasta"

def parse_fasta(file_like):
    """ generate that yields tuple containing record id, extra info
    in tail of header and the DNA sequence with newline characters
    """
    from Bio import SeqIO
    fasta_sequences = SeqIO.parse(file_like, 'fasta')
    for fasta in fasta_sequences:
        info, sequence = fasta.id, fasta.seq.tostring()
        # seprate record id from rest of the seqn info
        x = info.split('|')
        recID, recExtras = x[0], x[1:]
        yield recID, recExtras, sequence
    return

def prepare_data(txt_file, db_file):
    """ put data from txt_file into dct """
    i = 0
    with open(txt_file, mode='rt') as src_gen:
        fasta_gen = parse_fasta(src_gen)
        for recID, recExtras, dna_seq in fasta_gen:
            dct[recID] = dna_seq
            if i % 100 == 0:
                print(i, 'lines digested into sql database')
    return

def get_DNA_seq_of(recordID, src):
    """ search for recordID in src database and return a formatted string """
    ans = ""
    dna_seq = src[recordID]
    ans += ">%s\n%s\n" % (recordID, dna_seq)
    return ans

def main():
    # first of all prepare an optimized database
    dct = dict()
    prepare_data(OrigTaxonSeqsfile, dct)
    # now start searching and writing
    progress = 0
    out_file = open("02_Arthropoda_specimen_less.fasta", 'wt')
    taxa_file = open(Taxonomyfile, 'rt')
    with taxa_file, out_file:
        for line in taxa_file:
            question = line.split("\t")[0]
            answer = get_DNA_seq_of(question, dct)
            out_file.write(answer)
            if progress % 100 == 0:
                print(progress, 'lines processed')
    return

if __name__ == '__main__':
    main()

Answer 2

我认为你可以通过改进查找来创造巨大的性能提升。

使用set() 可以帮助您。集合旨在进行非常快速的数据查找，并且它们不存储重复值，这使它们成为过滤数据的理想选择。因此，让我们将输入文件中的所有分类 ID 存储在一个集合中。

from Bio import SeqIO
from Bio.Seq import Seq
import csv
import sys

taxonomy_file = "02_Arthropoda_specimen_data_less.txt"
orig_taxon_sequence_file = "00_Arthropoda_specimen.fasta"
output_sequence_file = "02_Arthropoda_specimen_less.fasta"

# build a set for fast look-up of IDs
with open(taxonomy_file, 'r', newline='') as fp:
    datareader = csv.reader(fp, delimiter='\t')
    first_column = (row[0] for row in datareader)
    taxonomy_ids = set(first_column)

# use the set to speed up filtering the input FASTA file
with open(output_sequence_file, 'w') as fp:
    for seq in SeqIO.parse(orig_taxon_sequence_file, "fasta"):
        if seq.id in taxonomy_ids: 
            fp.write('>')
            fp.write(seq.id)
            fp.write(seq.seq)
            fp.write('\n')

• 我已经重命名了一些变量。将变量命名为f4 只是在上面的评论中写入“#Name of the output sequence file”是完全没有意义的。为什么不去掉注释，直接将变量命名为output_sequence_file？
• (row[0] for row in datareader) 是一个生成器理解。生成器是一个可迭代对象，这意味着它还没有计算 ID 列表——它只知道要做什么。这通过不构建临时列表来节省时间和内存。一行之后，接受迭代的 set() 构造函数将使用第一列中的所有 ID 构建一个集合。
• 在第二个块中，我们使用if seq.id in taxonomy_ids 来检查是否应该输出序列ID。 in 在片场非常快。
• 我调用了四次.write()，而不是从四个项目中构建一个临时字符串。我假设 seq.id 和 seq.seq 已经是字符串，因此不需要对它们调用 str()。
• 我对 FASTA 文件格式了解不多，但快速浏览一下the BioPython documentation 表明使用SeqIO.write() 会是创建格式的更好方法。

Answer 3

我已在您的问题下的评论中要求澄清，但现在您没有回应（无意批评），所以我会在我必须离开之前尝试回答您的问题，基于以下假设我的代码。

1. 在第二个数据文件中，每条记录占用两行，第一行是排序标题，第二行是 ACGT 序列。
2. 在标题行中，我们有一个前缀 ">"，然后是一些由 "|" 分隔的字段，其中第一个字段是整个两行记录的 ID。

在上述假设下

# If possible, no hardcoded filenames, use sys.argv and the command line
import sys

# command line sanity check
if len(sys.argv) != 4:
     print('A descriptive error message')
     sys.exit(1)

# Names of the input and output files
fn1, fn2, fn3 = sys.argv[1:]

# Use a set comprehension to load the IDs from the first file
IDs = {line.split()[0] for line in open(fn1)} # a set

# Operate on the second file
with open(fn2) as f2:

    # It is possible to use `for line in f2: ...` but here we have(?)
    # two line records, so it's a bit different
    while True:

        # Try to read two lines from file
        try:
            header = f2.next()
            payload = f2.next()
        # no more lines? break out from the while loop...
        except StopIteration:
            break

        # Sanity check on the header line
        if header[0] != ">":
            print('Incorrect header line: "%s".'%header)
            sys.exit(1)

        # Split the header line on "|", find the current ID
        ID = header[1:].split("|")[0]

        # Check if the current ID was mentioned in the first file
        if ID in IDs:
            # your code

因为没有内部循环，这应该快 6 个数量级......它是否满足您的需要还有待观察 :-)