Split FASTQ(.gz) files into gzip chunks that stay under a size target.

split_fastq.py

#!/usr/bin/env python3
"""Split FASTQ(.gz) files into gzip chunks that stay under a size target.

This script is intended for customer-facing support workflows where a user needs
smaller `.fastq.gz` files for re-upload. It works on record boundaries and can
split paired-end R1/R2 inputs with identical record counts per chunk.

Because gzip compression ratio varies across reads, the script keeps a safety
margin below the requested maximum and rotates once the current chunk approaches
that threshold. The default settings target 200 MB output chunks with a 5 MB
safety buffer, which is conservative enough for typical Illumina FASTQ files.
"""

from __future__ import annotations

import argparse
import gzip
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO


DEFAULT_MAX_BYTES = 200_000_000
DEFAULT_SAFETY_BYTES = 10_000_000
DEFAULT_CHECK_EVERY = 10_000


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Split FASTQ(.gz) files into gzip chunks that stay below a target size."
        )
    )
    parser.add_argument(
        "--r1",
        required=True,
        type=Path,
        help="Path to the R1 FASTQ or FASTQ.GZ file.",
    )
    parser.add_argument(
        "--r2",
        type=Path,
        help="Optional path to the paired R2 FASTQ or FASTQ.GZ file.",
    )
    parser.add_argument(
        "--output-dir",
        required=True,
        type=Path,
        help="Directory where split files will be written.",
    )
    parser.add_argument(
        "--max-bytes",
        type=int,
        default=DEFAULT_MAX_BYTES,
        help=(
            f"Maximum desired output size per chunk in bytes. Default: "
            f"{DEFAULT_MAX_BYTES} (200 MB)."
        ),
    )
    parser.add_argument(
        "--safety-bytes",
        type=int,
        default=DEFAULT_SAFETY_BYTES,
        help=(
            f"Rotate before reaching this many bytes below --max-bytes. "
            f"Default: {DEFAULT_SAFETY_BYTES}."
        ),
    )
    parser.add_argument(
        "--check-every-records",
        type=int,
        default=DEFAULT_CHECK_EVERY,
        help=(
            "Flush outputs and evaluate chunk size every N records or read pairs. "
            f"Default: {DEFAULT_CHECK_EVERY}."
        ),
    )
    parser.add_argument(
        "--gzip-level",
        type=int,
        default=6,
        choices=range(1, 10),
        metavar="1-9",
        help="Gzip compression level for output chunks. Default: 6.",
    )
    return parser.parse_args()


def strip_fastq_suffix(path: Path) -> str:
    name = path.name
    for suffix in (".fastq.gz", ".fq.gz", ".fastq", ".fq"):
        if name.endswith(suffix):
            return name[: -len(suffix)]
    return path.stem


def open_fastq_reader(path: Path) -> BinaryIO:
    if path.suffix == ".gz":
        return gzip.open(path, "rb")
    return path.open("rb")


def read_fastq_record(handle: BinaryIO, label: str) -> bytes | None:
    header = handle.readline()
    if not header:
        return None

    sequence = handle.readline()
    plus = handle.readline()
    quality = handle.readline()

    if not sequence or not plus or not quality:
        raise ValueError(f"Incomplete FASTQ record encountered in {label}.")

    return header + sequence + plus + quality


@dataclass
class ChunkWriter:
    input_path: Path
    output_dir: Path
    gzip_level: int
    part_number: int = 0
    records_written: int = 0
    path: Path | None = None
    raw_handle: BinaryIO | None = None
    gzip_handle: gzip.GzipFile | None = None

    def open_next(self) -> None:
        self.close()
        self.part_number += 1
        self.records_written = 0
        filename = f"{strip_fastq_suffix(self.input_path)}.part{self.part_number:04d}.fastq.gz"
        self.path = self.output_dir / filename
        self.raw_handle = self.path.open("wb")
        self.gzip_handle = gzip.GzipFile(
            filename="",
            mode="wb",
            fileobj=self.raw_handle,
            compresslevel=self.gzip_level,
            mtime=0,
        )

    def write_record(self, record: bytes) -> None:
        if self.gzip_handle is None:
            self.open_next()
        assert self.gzip_handle is not None
        self.gzip_handle.write(record)
        self.records_written += 1

    def flush(self) -> None:
        if self.gzip_handle is not None:
            self.gzip_handle.flush()
        if self.raw_handle is not None:
            self.raw_handle.flush()

    def size_bytes(self) -> int:
        self.flush()
        if self.raw_handle is None:
            return 0
        return self.raw_handle.tell()

    def close(self) -> None:
        if self.gzip_handle is not None:
            self.gzip_handle.close()
            self.gzip_handle = None
        if self.raw_handle is not None:
            self.raw_handle.close()
            self.raw_handle = None


def rotate_needed(
    chunk_writers: list[ChunkWriter],
    threshold_bytes: int,
) -> bool:
    return any(writer.size_bytes() >= threshold_bytes for writer in chunk_writers)


def validate_args(args: argparse.Namespace) -> None:
    if args.max_bytes <= 0:
        raise ValueError("--max-bytes must be greater than 0.")
    if args.safety_bytes < 0:
        raise ValueError("--safety-bytes cannot be negative.")
    if args.safety_bytes >= args.max_bytes:
        raise ValueError("--safety-bytes must be smaller than --max-bytes.")
    if args.check_every_records <= 0:
        raise ValueError("--check-every-records must be greater than 0.")
    if not args.r1.exists():
        raise FileNotFoundError(f"R1 file not found: {args.r1}")
    if args.r2 and not args.r2.exists():
        raise FileNotFoundError(f"R2 file not found: {args.r2}")


def split_fastq_files(args: argparse.Namespace) -> int:
    output_dir: Path = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)

    for input_path in filter(None, [args.r1, args.r2]):
        stem = strip_fastq_suffix(input_path)
        for existing_path in output_dir.glob(f"{stem}.part*.fastq.gz"):
            existing_path.unlink()

    threshold_bytes = args.max_bytes - args.safety_bytes
    writers = [
        ChunkWriter(args.r1, output_dir, args.gzip_level),
    ]
    if args.r2:
        writers.append(ChunkWriter(args.r2, output_dir, args.gzip_level))

    total_units = 0

    with open_fastq_reader(args.r1) as r1_handle:
        r2_handle_cm = open_fastq_reader(args.r2) if args.r2 else None
        try:
            while True:
                r1_record = read_fastq_record(r1_handle, str(args.r1))
                if r1_record is None:
                    if r2_handle_cm is not None:
                        extra_r2 = read_fastq_record(r2_handle_cm, str(args.r2))
                        if extra_r2 is not None:
                            raise ValueError(
                                "R2 has more FASTQ records than R1. Inputs are not aligned."
                            )
                    break

                records_to_write = [r1_record]
                if r2_handle_cm is not None:
                    r2_record = read_fastq_record(r2_handle_cm, str(args.r2))
                    if r2_record is None:
                        raise ValueError(
                            "R1 has more FASTQ records than R2. Inputs are not aligned."
                        )
                    records_to_write.append(r2_record)

                if writers[0].gzip_handle is None:
                    for writer in writers:
                        writer.open_next()

                for writer, record in zip(writers, records_to_write, strict=True):
                    writer.write_record(record)

                total_units += 1

                if total_units % args.check_every_records == 0 and rotate_needed(
                    writers, threshold_bytes
                ):
                    for writer in writers:
                        writer.close()
        finally:
            if r2_handle_cm is not None:
                r2_handle_cm.close()
            for writer in writers:
                writer.close()

    return total_units


def summarize_outputs(output_dir: Path, writers: list[ChunkWriter]) -> list[str]:
    lines: list[str] = []
    for writer in writers:
        stem = strip_fastq_suffix(writer.input_path)
        chunk_paths = sorted(output_dir.glob(f"{stem}.part*.fastq.gz"))
        if not chunk_paths:
            continue
        lines.append(f"{writer.input_path.name}: {len(chunk_paths)} chunk(s)")
        for path in chunk_paths:
            size_bytes = path.stat().st_size
            lines.append(f"  {path.name}: {size_bytes} bytes")
    return lines


def find_oversized_outputs(
    output_dir: Path,
    writers: list[ChunkWriter],
    max_bytes: int,
) -> list[tuple[Path, int]]:
    oversized: list[tuple[Path, int]] = []
    for writer in writers:
        stem = strip_fastq_suffix(writer.input_path)
        for path in sorted(output_dir.glob(f"{stem}.part*.fastq.gz")):
            size_bytes = path.stat().st_size
            if size_bytes > max_bytes:
                oversized.append((path, size_bytes))
    return oversized


def main() -> int:
    args = parse_args()
    try:
        validate_args(args)
    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2

    writers = [ChunkWriter(args.r1, args.output_dir, args.gzip_level)]
    if args.r2:
        writers.append(ChunkWriter(args.r2, args.output_dir, args.gzip_level))

    try:
        total_units = split_fastq_files(args)
    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    print(
        f"Split {total_units:,} FASTQ "
        f"{'record pairs' if args.r2 else 'records'} into {args.output_dir}"
    )
    for line in summarize_outputs(args.output_dir, writers):
        print(line)

    oversized_outputs = find_oversized_outputs(args.output_dir, writers, args.max_bytes)
    if oversized_outputs:
        print(
            "error: one or more output chunks exceeded the requested size limit. "
            "Rerun with a larger --safety-bytes value or a smaller "
            "--check-every-records value.",
            file=sys.stderr,
        )
        for path, size_bytes in oversized_outputs:
            print(f"  {path.name}: {size_bytes} bytes", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())