perf(concat): block-copy concatenation instead of SeqRecord round-trip (89x faster, -2.6 GB)

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">")

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

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")