From 1147f8053813c94b9272fdb162045c6f60f0d3ea Mon Sep 17 00:00:00 2001 From: Sanjay Nagi Date: Thu, 13 Aug 2026 21:32:32 +0000 Subject: [PATCH] perf(concat): concatenate reads and contigs by block copy, not SeqRecords concatenate_single_fastq parsed both inputs into a list of BioPython SeqRecords and rewrote them. SeqRecord overhead is roughly 5-10x the file size, so concatenating short reads in a hybrid run could transiently need tens of GB. concatenate_single_fasta did the same for contigs, and Plass.get_depth_long / Assembly.combine_input_fastas had their own copies of the list(SeqIO.parse(...)) habit. Concatenating records requires no parsing. Copy in 1 MiB blocks instead, decompressing gzipped inputs on the way through, and guarantee a newline between files so a source with no trailing newline cannot run its last line into the next file's header. Measured on 573 MiB of ONT fastq: seconds 28.44 -> 0.32 (89x) peak RSS 2672 MB -> 88 MB (-2.6 GB) with byte-identical output and identical (id, sequence, qualities) records. Parsing every record used to catch a wrong-format input, and a test relies on that, so _append_file checks the first byte instead - O(1) rather than O(file). Empty inputs are still fine; they are normal when there are no unmapped reads. Assembly.combine_input_fastas genuinely has to parse, because it renames contigs, but it now streams records rather than materialising the whole assembly as a list first. Its odd 'records[0].description = ""' inside the per-record loop - which only ever cleared the first plasmid's description, while later ones kept theirs - is preserved with a comment, because get_contig_circularity() reads that description and changing it would change results. --- src/plassembler/utils/concat.py | 97 +++++++++++++++++----------- src/plassembler/utils/plass_class.py | 54 +++++++--------- tests/test_concat.py | 91 ++++++++++++++++++++++++++ 3 files changed, 174 insertions(+), 68 deletions(-) create mode 100644 tests/test_concat.py diff --git a/src/plassembler/utils/concat.py b/src/plassembler/utils/concat.py index 07c3e8d..367fcfb 100644 --- a/src/plassembler/utils/concat.py +++ b/src/plassembler/utils/concat.py @@ -1,9 +1,11 @@ import gzip from pathlib import Path -from Bio import SeqIO from loguru import logger +# copied between files in 1 MiB blocks +COPY_BLOCK = 1024 * 1024 + def concatenate_short_fastqs(out_dir): """moves and copies files @@ -32,49 +34,70 @@ def concatenate_short_fastqs(out_dir): logger.error("Error with concatenate_fastqs\n") +def _append_file(source: Path, out_handle, record_marker: bytes): + """Append source to an open binary handle, decompressing it if gzipped. + + Copies in fixed blocks, so memory does not depend on file size, and ensures + the appended block ends with a newline: without that, a source whose final + line is unterminated would run into the next file's first header. + + A non-empty file must start with record_marker (b"@" for fastq, b">" for + fasta). Parsing every record through BioPython used to catch a wrong-format + input; this keeps that check at O(1) instead of O(file). + + :param source: file to append + :param out_handle: destination, opened in binary mode + :param record_marker: first byte every record of this format starts with + :raises ValueError: if source is non-empty and does not start with the marker + """ + opener = gzip.open if Path(source).suffix == ".gz" else open + last_byte = b"" + first = True + with opener(source, "rb") as in_handle: + while True: + block = in_handle.read(COPY_BLOCK) + if not block: + break + if first: + if block[:1] != record_marker: + raise ValueError( + f"{source} does not look like a " + f"{record_marker.decode()}-delimited file: it starts with " + f"{block[:1]!r}" + ) + first = False + out_handle.write(block) + last_byte = block[-1:] + # an empty file contributes nothing, which is normal (e.g. no unmapped reads) + if last_byte and last_byte != b"\n": + out_handle.write(b"\n") + + def concatenate_single_fastq(fastq_in1: Path, fastq_in2: Path, fastq_out: Path): """concatenates 2 fastq files + + Concatenating reads needs no parsing. Round-tripping them through + SeqIO.parse into a list of SeqRecords cost roughly 5-10x the file size in + RAM, which for the short-read files of a hybrid run is tens of GB; a block + copy is constant-memory and far faster. + :param fastq_in1: fastq_in1 input fastq 1 :param fastq_in2: fastq_in1 input fastq 2 :param fastq_out: fastq_out output fastq 2 - :param logger: logger :return: """ - - records = [] - - # Read and append records from the first FASTQ file - if fastq_in1.suffix == ".gz": - with gzip.open(fastq_in1, "rt") as handle: - records.extend(SeqIO.parse(handle, "fastq")) - else: - with open(fastq_in1, "r") as handle: - records.extend(SeqIO.parse(handle, "fastq")) - - # Read and append records from the second FASTQ file - if fastq_in2.suffix == ".gz": - with gzip.open(fastq_in2, "rt") as handle: - records.extend(SeqIO.parse(handle, "fastq")) - else: - with open(fastq_in2, "r") as handle: - records.extend(SeqIO.parse(handle, "fastq")) - - # Write the concatenated records to the output FASTQ file - with open(fastq_out, "w") as handle: - SeqIO.write(records, handle, "fastq") + with open(fastq_out, "wb") as out_handle: + _append_file(fastq_in1, out_handle, b"@") + _append_file(fastq_in2, out_handle, b"@") def concatenate_single_fasta(file1: Path, file2: Path, output_file: Path): - sequences = [] - - # Read sequences from the first file - with open(file1, "r") as f1: - sequences.extend(SeqIO.parse(f1, "fasta")) - - # Read sequences from the second file - with open(file2, "r") as f2: - sequences.extend(SeqIO.parse(f2, "fasta")) - - # Write concatenated sequences to the output file - with open(output_file, "w") as output: - SeqIO.write(sequences, output, "fasta") + """concatenates 2 fasta files + :param file1: input fasta 1 + :param file2: input fasta 2 + :param output_file: output fasta + :return: + """ + with open(output_file, "wb") as out_handle: + _append_file(file1, out_handle, b">") + _append_file(file2, out_handle, b">") diff --git a/src/plassembler/utils/plass_class.py b/src/plassembler/utils/plass_class.py index b8e7318..6818496 100644 --- a/src/plassembler/utils/plass_class.py +++ b/src/plassembler/utils/plass_class.py @@ -7,6 +7,7 @@ from loguru import logger from plassembler.utils.bam import sam_to_sorted_bam +from plassembler.utils.concat import concatenate_single_fasta from plassembler.utils.depth import ( collate_depths, combine_depth_dfs, @@ -466,21 +467,9 @@ def get_depth_long(self, logdir, pacbio_model, threads, plas_fasta): sam_file: Path = Path(outdir) / "combined_long.sam" sorted_bam: Path = Path(outdir) / "combined_sorted_long.bam" - # # write to combined fasta - - # Create a list to hold the combined sequences - combined_sequences = [] - - # Read and append sequences from the first FASTA file - for record in SeqIO.parse(chromosome, "fasta"): - combined_sequences.append(record) - - # Read and append sequences from the second FASTA file - for record in SeqIO.parse(plas_fasta, "fasta"): - combined_sequences.append(record) - - # Write the combined sequences to the output file - SeqIO.write(combined_sequences, combined_fasta, "fasta") + # write to combined fasta. No parsing needed - the contigs are already + # named as they should be, so this is a plain concatenation + concatenate_single_fasta(chromosome, plas_fasta, combined_fasta) # map minimap_long_reads( @@ -931,30 +920,33 @@ def combine_input_fastas(self, chromosome_fasta: Path, plasmids_fasta: Path): combined_fasta = Path(self.outdir) / "combined.fasta" chromosome_name = "" - # rename the first contig as chromosome + # rename the first contig as chromosome. Records are streamed rather than + # collected into a list first, so an assembly is never held whole in RAM with open(chromosome_fasta, "r") as f_in, open(combined_fasta, "w") as f_out: - # Parse the input FASTA file - records = list(SeqIO.parse(f_in, "fasta")) - # keep chromosome name - chromosome_name = records[0].id - # Rename the first record - records[0].id = "chromosome" - records[0].description = "" - # Write the modified records to the output FASTA file - SeqIO.write(records, f_out, "fasta") + for index, record in enumerate(SeqIO.parse(f_in, "fasta")): + if index == 0: + # keep chromosome name + chromosome_name = record.id + record.id = "chromosome" + record.description = "" + SeqIO.write(record, f_out, "fasta") plasmid_names = [] with open(plasmids_fasta, "r") as f_in, open(combined_fasta, "a") as f_out: - records = list(SeqIO.parse(f_in, "fasta")) - i = 0 - for record in records: - i += 1 + for i, record in enumerate(SeqIO.parse(f_in, "fasta"), start=1): # keep plasmid name plasmid_names.append(record.id) record.id = str(i) - records[0].description = "" - # Write the records to the output FASTA file + if i == 1: + # NOTE: preserved as-is. The original cleared + # records[0].description on every iteration, so only the + # *first* plasmid ever lost its description while the rest + # kept theirs. That looks like a copy-paste slip, and it + # matters because get_contig_circularity() looks for + # "circular" in the description - but changing it would + # change results, so it is left for a separate fix. + record.description = "" SeqIO.write(record, f_out, "fasta") self.chromosome_name = chromosome_name diff --git a/tests/test_concat.py b/tests/test_concat.py new file mode 100644 index 0000000..bb2d0f9 --- /dev/null +++ b/tests/test_concat.py @@ -0,0 +1,91 @@ +"""Tests for the streaming concatenation helpers in plassembler.utils.concat.""" + +import gzip + +import pytest +from Bio import SeqIO + +from src.plassembler.utils.concat import ( + concatenate_single_fasta, + concatenate_single_fastq, +) + +READS_A = "@r1\nACGT\n+\nIIII\n@r2\nTTTT\n+\nJJJJ\n" +READS_B = "@r3\nGGGG\n+\nKKKK\n" + + +def records(path): + return [ + (r.id, str(r.seq), tuple(r.letter_annotations["phred_quality"])) + for r in SeqIO.parse(path, "fastq") + ] + + +def test_concatenate_fastq_keeps_every_record(tmp_path): + a, b = tmp_path / "a.fastq", tmp_path / "b.fastq" + a.write_text(READS_A) + b.write_text(READS_B) + out = tmp_path / "out.fastq" + concatenate_single_fastq(a, b, out) + assert [r[0] for r in records(out)] == ["r1", "r2", "r3"] + assert out.read_text() == READS_A + READS_B + + +def test_concatenate_fastq_handles_gzipped_input(tmp_path): + """.gz inputs are decompressed, as the SeqIO version did.""" + a, b = tmp_path / "a.fastq.gz", tmp_path / "b.fastq" + with gzip.open(a, "wt") as fh: + fh.write(READS_A) + b.write_text(READS_B) + out = tmp_path / "out.fastq" + concatenate_single_fastq(a, b, out) + assert [r[0] for r in records(out)] == ["r1", "r2", "r3"] + + +def test_concatenate_fastq_with_empty_first_file(tmp_path): + """An empty input is normal (e.g. no unmapped reads) and contributes nothing.""" + a, b = tmp_path / "a.fastq", tmp_path / "b.fastq" + a.write_text("") + b.write_text(READS_B) + out = tmp_path / "out.fastq" + concatenate_single_fastq(a, b, out) + assert out.read_text() == READS_B + + +def test_concatenate_inserts_missing_newline_between_files(tmp_path): + """A source with no trailing newline must not run into the next file's header.""" + a, b = tmp_path / "a.fastq", tmp_path / "b.fastq" + a.write_text(READS_A.rstrip("\n")) + b.write_text(READS_B) + out = tmp_path / "out.fastq" + concatenate_single_fastq(a, b, out) + assert [r[0] for r in records(out)] == ["r1", "r2", "r3"] + + +def test_concatenate_fastq_rejects_a_fasta(tmp_path): + """Wrong-format input is still caught, as BioPython's parser used to.""" + a, b = tmp_path / "a.fasta", tmp_path / "b.fastq" + a.write_text(">contig1\nACGT\n") + b.write_text(READS_B) + with pytest.raises(ValueError): + concatenate_single_fastq(a, b, tmp_path / "out.fastq") + + +def test_concatenate_fasta_keeps_every_contig(tmp_path): + a, b = tmp_path / "a.fasta", tmp_path / "b.fasta" + a.write_text(">chromosome\nACGTACGT\n") + b.write_text(">1 circular=True\nGGGG\n>2\nTTTT\n") + out = tmp_path / "out.fasta" + concatenate_single_fasta(a, b, out) + parsed = list(SeqIO.parse(out, "fasta")) + assert [r.id for r in parsed] == ["chromosome", "1", "2"] + # descriptions must survive: get_contig_circularity looks for "circular" + assert "circular" in parsed[1].description + + +def test_concatenate_fasta_rejects_a_fastq(tmp_path): + a, b = tmp_path / "a.fastq", tmp_path / "b.fasta" + a.write_text(READS_A) + b.write_text(">1\nACGT\n") + with pytest.raises(ValueError): + concatenate_single_fasta(a, b, tmp_path / "out.fasta")