From 3e580ad9e4eaed800d624ce966236aaa365402a6 Mon Sep 17 00:00:00 2001 From: gwct Date: Thu, 10 Sep 2026 14:06:40 -0400 Subject: [PATCH] Stream indexes in fetch/gc/stats so whole-genome MAFs dont exhaust memory --- DEVELOPMENT.md | 150 ++++++++- mafutils/fetch.py | 563 +++++++++++++++++++++++++++------ mafutils/gc.py | 114 ++++--- mafutils/lib/common.py | 51 +++ mafutils/stats.py | 252 ++++++++++----- tests/example.maf.scaffold.idx | 13 +- tests/test_compression.py | 1 - tests/test_fetch.py | 51 +-- tests/test_fetch_streaming.py | 177 +++++++++++ tests/test_scaffold_mode.py | 116 +++++++ tests/test_stats.py | 23 ++ tests/test_validate.py | 38 ++- 12 files changed, 1287 insertions(+), 262 deletions(-) create mode 100644 tests/test_fetch_streaming.py create mode 100644 tests/test_scaffold_mode.py diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 0a2cf89..93e639b 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -4,15 +4,15 @@ Internal implementation notes, gotchas, and maintainer workflows (testing, releasing) for developing `mafutils` itself — see [`README.md`](README.md) for user-facing installation/usage docs. -- `tests/example.maf.scaffold.idx` is preserved as an older, headerless - scaffold-index fixture for comparison — it deliberately sits at - `example.maf`'s *default* scaffold-index path - (`common.deriveScaffoldIndexPath`), so `mafutils validate` against the - checked-in `example.maf`/`example.maf.block.idx` pair reports - UNVERIFIABLE rather than VERIFIED (it can't cross-check a headerless - scaffold index against the block index — see below). Tests exercising - the "clean, fully verified match" case build a fresh pair with - `mafutils index` in a tmp dir instead of using this fixture pair directly. +- `tests/example.maf.scaffold.idx` is a **current-format** index, consistent + with `example.maf.block.idx`. It used to be a deliberately preserved + headerless/off-by-one fixture, but `fetch` block mode now reads it to order + regions (see the streaming note below), and a stale pair silently drops + regions — an off-by-one run boundary is enough to skip every run's first + block. Tests that need a *headerless* index now build one in `tmp_path` by + stripping the header (`writeHeaderlessIndex` in `tests/test_validate.py`), + so no committed file's staleness is load-bearing any more. Keep these two + index files in sync: regenerate both with `mafutils index` together. - `tests/example.maf.scaffold.regenerated.idx` is the scaffold index produced by the current `mafutils index` implementation. - The real production scaffold indexes checked so far match the regenerated @@ -54,8 +54,8 @@ for user-facing installation/usage docs. is backward compatible. Format-1 indexes (pre-dating size/mtime/hash) simply lack those keys in the parsed header dict — every consumer treats their absence as "can't check this dimension," not an error, same as a - fully missing header. `tests/example.maf.scaffold.idx` is deliberately - kept in format-1 shape (see above) specifically to exercise this path. + fully missing header. Tests build a headerless index on the fly rather than + relying on a committed one (see above). - **Computing the hash costs nothing extra at `mafutils index` time, for all three compression types** — this was the whole point of always storing it. `common._HashingRawIO` (an `io.RawIOBase` subclass) sits @@ -168,6 +168,134 @@ for user-facing installation/usage docs. **bgzip deliberately stays on the text+`tell()` path** — its virtual offsets are not additive, so a byte counter cannot reproduce them (same root cause as the straddling-`a`-line bug above). +- **Whole-scaffold extraction must stream; `readMafBlockBytes` is only for + block-sized ranges.** `fetch -m scaffold` used to do + `readMafBlockBytes(...)` for a whole scaffold, which materializes the range + as one bytes object. On a real 7.4TB 241-species MAF, chr1 spans bytes + 16 -> 597,495,640,164, so that became a single **~597GB `read()`** and died + with a `MemoryError`. Two traps worth knowing: `str(MemoryError())` is the + **empty string**, so the failure logged as a completely blank + `[ERROR] chr1:` message; and with memory overcommit the doomed read + actually starts, so at `-p 1` it looks like a hang rather than an error. + `common.copyMafRangeToStream()` now copies in bounded chunks (O(chunk_size) + memory regardless of range size) straight into the output handle. Note its + bgzip branch also had to change: the old line-by-line loop accumulated + lines into a list to `join()` at the end, so it was O(range) too -- + fixing only the uncompressed branch would have left bgzip broken at scale. +- **Scaffold extraction failures now fail the run.** Both `fetchByScaffold` + and `fetchScaffoldsSequential` used to catch any exception, log it, and + then report `"Wrote "` regardless, so `fetch` exited **0** with a + 0-byte output file -- a downstream Snakemake pipeline consumed that as + success and only broke several steps later. They now return a status dict, + the caller `sys.exit(1)`s if any scaffold failed *or* extracted 0 bytes, + "Wrote" is logged only for genuinely non-empty output, and errors are + logged with `repr(e)` plus a debug traceback (precisely because the + exception that actually occurred had an empty `str()`). This was also a + direct violation of `AGENTS.md`'s "never insert silent fallbacks or error + handling that could hide errors" rule -- worth re-reading the other + `except Exception` sites against that standard. +- **`gc` streams the block index rather than loading it.** `gc.parseIndex` + built a list of every block's `(offset_start, offset_end)`, then + `list(chunker(...))` **sliced that list**, making a second full copy. + Measured peak, scaling perfectly linearly: 34MB @ 250k blocks, 136MB @ 1M, + 563MB @ 4M -- extrapolating to **~32GB at the 233M blocks** of a real + whole-genome index (observed as a 62GB-RSS OOM kill before any GC was + computed). Replaced by `iterIndexEntries()` (a generator) plus an + `islice`-based `chunker()` that never materializes its source, with a + bounded number of in-flight futures and results merged as they arrive via + `mergeCountsInto()` -- submitting all ~47k chunks upfront, or collecting + every worker's count dict before merging, would each have reintroduced the + same scaling. Peak is now **flat at ~1.5MB** regardless of index size. +- **`stats` streams its index the same way**, with three wrinkles `gc` didn't + have. (1) `block_id` is assigned during index parsing, and the original + incremented it *before* the `int()` conversions, so a line with enough + fields that fails to parse still consumes an id; `iterIndexEntries` + reproduces that exactly, since `block_id` is written into the per-block + output table. (2) The gz path used to set `chunk_size = len(entries)` to + force one single task, keeping that one worker's seeks forward-only -- + which requires knowing the entry count upfront, the one thing streaming + can't do. `workerTask` now takes an optional `maf_fp`, and the + single-process path opens one handle and reuses it across chunks, which + preserves forward-only seeking without materializing anything (measured: + no gz slowdown, 41.3s -> 40.4s). (3) Results are merged on arrival via + `mergeStatsInto`/`newStatsAccumulator` rather than collected -- ~47k chunk + results, each carrying a per-species dict, would otherwise scale with + index size just like the entries list did. +- **`stats`' per-block output rows were being concatenated in lexicographic + filename order** -- a pre-existing bug, unrelated to streaming, found while + testing the above. Each chunk writes `maf_stats.blocks.task.tsv` and + `writeBlockTable` did `sorted(block_tmp_paths)`, so `task10` sorted before + `task2` and `block.tsv` came out badly out of genomic order for any run + with >=10 chunks (a real 269k-block run has 54). Results are now tracked + as `(task_id, path)` and sorted numerically; verified 0 out-of-order + transitions across 269,308 rows. Note this *changes* `block.tsv` row order + versus older mafutils for multi-chunk runs -- it was wrong before. +- **`fetch` block mode streams the index too, by sorting regions into index + order instead of doing random access.** It used to build a 5-key **dict per + block** (~288 bytes/block, ~67GB at 233M blocks) so it could `bisect` by + coordinate. Now regions are sorted into index order and each worker walks its + own span of the index forward, matching regions as it passes them -- regions + are the cheap side (979k regions vs 233M blocks, ~240x fewer). Measured on + 20k real regions against the 8.09M-block hamster index: **3.30GB -> 47MB peak + RSS and 8.8x faster** (68.7s -> 7.7s; the old path spent most of its time + building that structure and pickling it to every worker), with all 20,000 + per-region outputs byte-identical for `none`/`gz`/`bgzip`. + Four things that make this correct, each of which broke a first attempt: + 1. **Scaffold order is MAF file order, not alphabetical.** Real data runs + `CM000994.3, GL456210.1, ..., CM000995.3`. The ordering comes from + `.scaffold.idx`, whose line order is exactly the block index's. + Sorting by scaffold *name* would seek backwards through the index. + 2. **A scaffold can occupy several disjoint runs.** MAFs may interleave + scaffolds -- `tests/example.maf` does (`chr4 ... chrX ... chr4`), so + `chr4` has two runs. `readScaffoldRuns` returns one entry per *run*, not + per scaffold; assuming one contiguous run per scaffold silently loses the + later blocks. + 3. **Runs are bounded by their MAF byte range, not by index byte offsets.** + `findIndexOffsetForMafByte` binary-searches the index *by byte position* + (seek to a midpoint, discard the partial line, compare the next complete + record -- ~35 probes for 35GB, the `look(1)` technique) using + `offset_start`, the only column that increases monotonically through the + whole file. Its result is a valid place to *start* scanning, deliberately + allowed to sit early, so it must never be treated as an exact boundary. + 4. **Offsets passed between helpers are exact record boundaries.** + `iterIndexRecordsFrom` does NOT skip a leading partial line, because + resume offsets point at real record starts; skipping one dropped the first + record of every region after the first. + Entries reach `fetchByRegion` as an **iterator**, never a list: one BED region + can span a whole scaffold, so materializing its blocks would reintroduce an + O(blocks) blowup for a single region. (FASTA output is still inherently + O(region length x species), since `fasta_seqs` accumulates before writing.) + Block mode now also requires the block and scaffold indexes to come from the + same `mafutils index` run, and errors if their headers disagree -- it locates + blocks via byte ranges recorded in the scaffold index, so a mismatched pair + silently drops regions. +- **`prefetchBlockCache` is gone.** It existed to give gzip "read every needed + block once, in ascending file order", but it materialized the decoded text of + every needed block (a real run logged "Prefetching 1031475 distinct + block(s)") -- its own O(needed blocks) memory blowup. Streaming the index in + order gives that property to *all three* compression types structurally, and + reuse between neighbouring regions is covered by the bounded + `WORKER_BLOCK_CACHE` LRU, since sorting makes such regions adjacent. +- **`fetch`'s no-output accounting read the wrong column, and it mattered a + lot.** The per-region summary line is `scaffold, start, end, basename, + n.overlapping.blocks, ...`; the code checked field **1** (`start`) instead of + field **4** (`n.overlapping.blocks`), so it was really asking "does this + region start at coordinate 0?". Two opposite symptoms: a region starting at 0 + that worked fine was counted as "no output" (a one-region BED at coordinate 0 + wrote correct output and then **failed the run**), while a region elsewhere + that found zero blocks was counted as written. On `tests/example.bed` the two + errors cancelled to exactly 1/10 = 0.10 -- precisely the threshold, which + fires on `>` -- so every test passed and the thresholds were never really + exercised. +- **Non-overlap outcomes are advisory, not fatal.** With the count fixed, the + leftover hardcoded thresholds from the abandoned `--max-no-overlap-*` feature + started failing legitimate runs (`example.bed` genuinely has 2/10 = 20% + non-overlapping). `fetch` is normally pointed at many intervals, so + `ADVISORY_NO_OVERLAP_REGIONS`/`ADVISORY_NO_OVERLAP_FRACTION`, "all regions + produced nothing", and "BED names a scaffold absent from the index" all now + **warn and continue**. Real misuse is still caught loudly and earlier: + index/MAF size+hash mismatch, a mismatched block/scaffold index pair, and + per-scaffold extraction failures. - The vendored `bgzf.BgzfReader` has no `.name` attribute (unlike `gzip.GzipFile`), so `fetch.py` gets the MAF's display filename from the known `maf_file` path (`WORKER_MAF_FILE`), not from the open file handle. diff --git a/mafutils/fetch.py b/mafutils/fetch.py index a15ce78..8116148 100644 --- a/mafutils/fetch.py +++ b/mafutils/fetch.py @@ -67,6 +67,7 @@ import logging import bisect import time +import traceback from enum import Enum from types import SimpleNamespace from typing import Annotated, Optional @@ -80,14 +81,16 @@ ############################################################################# -MAX_NO_OVERLAP_REGIONS_DEFAULT = 10000 -MAX_NO_OVERLAP_FRACTION_DEFAULT = 0.10 +ADVISORY_NO_OVERLAP_REGIONS = 10000 +ADVISORY_NO_OVERLAP_FRACTION = 0.10 WORKER_BLOCK_CACHE_MAX = 2048 WORKER_HEADER = None WORKER_MAF_FILE = None WORKER_MAF_COMPRESSION = None WORKER_INDEX = None +WORKER_INDEX_FILE = None +WORKER_SCAFFOLD_RUNS = None WORKER_OUTPUT = None WORKER_SINGLE_OUTPUT = None WORKER_SCAFFOLD_SUBDIRS = None @@ -176,11 +179,15 @@ def initBatchWorker( profile, prefetched_cache=None, scaffold_subdirs=False, + index_file=None, + scaffold_runs=None, ): global WORKER_HEADER global WORKER_MAF_FILE global WORKER_MAF_COMPRESSION global WORKER_INDEX + global WORKER_INDEX_FILE + global WORKER_SCAFFOLD_RUNS global WORKER_OUTPUT global WORKER_SINGLE_OUTPUT global WORKER_SCAFFOLD_SUBDIRS @@ -197,6 +204,11 @@ def initBatchWorker( WORKER_MAF_FILE = maf_file WORKER_MAF_COMPRESSION = maf_compression WORKER_INDEX = index + # Block mode streams the index from disk rather than receiving it: the + # in-memory form is ~288 bytes/block (~67GB for a whole-genome index), and + # passing it through initargs pickled that whole structure to EVERY worker. + WORKER_INDEX_FILE = index_file + WORKER_SCAFFOLD_RUNS = scaffold_runs WORKER_OUTPUT = output WORKER_SINGLE_OUTPUT = single_output WORKER_SCAFFOLD_SUBDIRS = scaffold_subdirs @@ -292,6 +304,138 @@ def parseIndex(index_file, LOG, mode="block"): ############################################################################# +def readScaffoldRuns(scaffold_index_file): + """ + Returns the scaffold index's rows in file order as + [(scaffold, maf_start_byte, maf_end_byte), ...]. + + Each row is a *run* of consecutive blocks for one scaffold, not a unique + scaffold: a MAF may interleave scaffolds, so the same name can appear more + than once. tests/example.maf does exactly this (chr4 ... chrX ... chr4), + and assuming one run per scaffold silently loses the later blocks. The + block index is in MAF order too, so these runs are also contiguous there. + """ + runs = [] + with open(scaffold_index_file, "r", encoding="utf-8") as fp: + for line in fp: + if line.startswith("#") or not line.strip(): + continue + fields = line.rstrip("\n").split("\t") + if len(fields) < 3: + continue + try: + runs.append((fields[0], int(fields[1]), int(fields[2]))) + except ValueError: + continue + return runs + + +def parseIndexLine(line): + """ + Parses one block-index line into + (scaffold, ref_start, ref_length, offset_start, offset_end), or None if the + line is a comment/blank/malformed. Only the columns fetch actually needs. + """ + if not line or line.startswith("#") or not line.strip(): + return None + fields = line.rstrip("\n").split("\t") + if len(fields) < 8: + return None + try: + return (fields[0], int(fields[1]), int(fields[2]), int(fields[6]), int(fields[7])) + except ValueError: + return None + + +def _firstRecordAtOrAfter(fp, offset): + """ + Seeks to `offset` and returns (record, record_offset) for the first + COMPLETE record at or after it, or (None, None) at EOF. + + Index lines are variable-length, so there is no arithmetic from "record N" + to a byte position -- but seeking to an arbitrary byte and discarding the + partial line lands on a record boundary, which is what makes a binary + search over the text file possible. + """ + fp.seek(offset) + if offset > 0: + fp.readline() # discard the (probably partial) line we landed inside + while True: + record_offset = fp.tell() + line = fp.readline() + if not line: + return None, None + record = parseIndexLine(line) + if record is not None: + return record, record_offset + + +def findIndexOffsetForMafByte(index_file, maf_offset, window=8192): + """ + Byte offset into the block index of the first record whose block starts at + or after `maf_offset` in the MAF. + + The search key is the block's MAF offset because that is the one column + guaranteed to increase monotonically through the whole index -- scaffold + name and ref_start are only sorted *within* a run, and a scaffold can have + several runs (see readScaffoldRuns). + + Binary search by byte position: seek to a midpoint, skip the partial line, + compare the first complete record, narrow. ~log2(filesize) probes, so ~35 + on a 35GB index regardless of how many records it holds. Returns an exact + record boundary, since callers treat offsets as record starts. + """ + size = os.path.getsize(index_file) + with open(index_file, "r", encoding="utf-8") as fp: + lo, hi = 0, size + while hi - lo > window: + mid = (lo + hi) // 2 + record, _ = _firstRecordAtOrAfter(fp, mid) + if record is None: + hi = mid + continue + if record[3] < maf_offset: + lo = mid + else: + hi = mid + _record, record_offset = _firstRecordAtOrAfter(fp, lo) + return size if record_offset is None else record_offset + + +def computeRunIndexOffsets(index_file, runs): + """ + Maps each scaffold-index run to where its records begin in the block index. + + One binary search per run (61 runs on the hamster data, 455 on the + 241-mammal file) done once at startup -- not per region. Run i's records + occupy [offsets[i], offsets[i+1]). + """ + return [findIndexOffsetForMafByte(index_file, maf_start) for _scaffold, maf_start, _maf_end in runs] + + +def iterIndexRecordsFrom(fp, offset): + """ + Yields (record, record_offset) from `offset` forward on an already-open + index handle. Streaming, so memory never scales with index size. + + `offset` MUST be an exact record boundary (both findIndexOffsetForMafByte + and the record_offset values yielded here are). It deliberately does not + skip a leading partial line -- doing so would drop the first record when + resuming from a known boundary. + """ + fp.seek(offset) + while True: + record_offset = fp.tell() + line = fp.readline() + if not line: + return + record = parseIndexLine(line) + if record is not None: + yield record, record_offset + + +############################################################################# + def parseBed(bed_file, LOG, mode="block"): """ Parses the BED file into a list of region dictionaries. @@ -462,6 +606,59 @@ def findOverlappingEntries(scaffold, bed_start, bed_end, index, profile_state=No return overlapping +def iterRegionEntries(index_fp, scaffold, bed_start, bed_end, run_indices, runs, run_offsets, resume, state): + """ + Streams the blocks overlapping `scaffold`:[bed_start, bed_end), walking each + run of that scaffold in file order. + + Runs are bounded by their MAF byte range (from the scaffold index), NOT by + index byte positions: findIndexOffsetForMafByte returns a valid place to + *start* scanning, deliberately allowed to sit early, so it must never be + treated as an exact boundary. Each record's own offset_start is compared + against the run's MAF range instead, which is exact. + + `resume` maps run -> byte offset to restart from, so consecutive regions on + a scaffold do not rescan a run from its beginning. Because a worker walks + its regions in ascending coordinate order, every later region has + start >= bed_start, so any record ending at or before bed_start is + irrelevant to all of them -- making the first record ending after bed_start + a safe resume point, recorded in state. + + Nothing is buffered: entries are yielded as they stream, so a region + spanning an entire scaffold costs O(1) memory here rather than O(blocks). + """ + new_resume = dict(resume) + for run in run_indices: + _run_scaffold, maf_start, maf_end = runs[run] + scan_from = max(run_offsets[run], resume.get(run, 0)) + resume_set = False + + for record, record_offset in iterIndexRecordsFrom(index_fp, scan_from): + rec_scaffold, ref_start, ref_length, offset_start, offset_end = record + if offset_start >= maf_end: + break # past this run in the MAF + if offset_start < maf_start: + continue # not yet into this run + if rec_scaffold != scaffold: + continue # defensive; a run is single-scaffold + if ref_start + ref_length <= bed_start: + continue # block ends before the region starts + if not resume_set: + new_resume[run] = record_offset + resume_set = True + if ref_start >= bed_end: + break # ascending, so nothing further overlaps + yield { + "ref_start": ref_start, + "ref_length": ref_length, + "offset_start": offset_start, + "offset_end": offset_end, + } + state["resume"] = new_resume + + +############################################################################# + def prefetchBlockCache(regions, index, maf_file, maf_compression, LOG): """ For plain-gzip input: collects every block entry needed across ALL @@ -736,7 +933,7 @@ def fetchByRegion( region, header, maf_fp, - index, + overlapping_entries, output, BATCHLOG, single_output=False, @@ -788,7 +985,9 @@ def fetchByRegion( blocks_written = 0; - if scaffold not in index: + if overlapping_entries is None: + # Caller signals "this scaffold isn't in the index at all" by passing + # None, rather than fetchByRegion holding an index to check against. appendWarning( warning_state, "missing-scaffold", @@ -800,8 +999,13 @@ def fetchByRegion( blocks_written = 0 scan_timer_start = time.perf_counter() if profile_state is not None else None - overlapping_entries = findOverlappingEntries(scaffold, bed_start, bed_end, index, profile_state=profile_state) - + # overlapping_entries is an ITERATOR, deliberately not a list: a single BED + # region can span an entire scaffold (hundreds of millions of blocks on a + # whole-genome MAF), so materializing its blocks would reintroduce an + # O(blocks) blowup for one region. Trimmed output is written to out_stream + # inside this loop, so MAF output stays O(1) in block count. (FASTA output + # is inherently O(region length x species) -- fasta_seqs below accumulates + # each species' sequence before writing.) for entry in overlapping_entries: try: read_timer_start = time.perf_counter() if profile_state is not None else None @@ -986,49 +1190,86 @@ def fetchByBatch( profile_state = initProfileState() if WORKER_PROFILE else None batch_timer_start = time.perf_counter() if WORKER_PROFILE else None - ordered_regions = sorted( - enumerate(batch), - key=lambda item: ( - item[1]["scaffold"], - item[1]["start"], - item[1]["end"], - item[0], - ), - ) + # Regions arrive already sorted into INDEX order by the parent, which is + # what lets this walk the index forward once. Re-sorting here (as this used + # to, by scaffold *name*) would break that: index scaffold order is MAF file + # order, not alphabetical. + index_fp = open(WORKER_INDEX_FILE, "r", encoding="utf-8") + scaffold_to_runs, runs, run_offsets = WORKER_SCAFFOLD_RUNS + # Per-run resume points, so consecutive regions on a scaffold continue + # forward through the index instead of rescanning from the run's start. + resume_by_scaffold = {} - for batch_order, (result_idx, region) in enumerate(ordered_regions, start=1): - region_num += 1 + try: + for batch_order, (result_idx, region) in enumerate(enumerate(batch), start=1): + region_num += 1 - batch_str = f" [batch {batch_num}.{batch_order}]" - BATCHLOG.debug(f">>>{batch_str} Region {region['scaffold']}:{region['start']}-{region['end']}") + batch_str = f" [batch {batch_num}.{batch_order}]" + BATCHLOG.debug(f">>>{batch_str} Region {region['scaffold']}:{region['start']}-{region['end']}") + + scaffold = region["scaffold"] + run_indices = scaffold_to_runs.get(scaffold) + if not run_indices: + entries = None # signals "scaffold absent from the index" + scan_state = None + else: + scan_state = {} + entries = iterRegionEntries( + index_fp, + scaffold, + region["start"], + region["end"], + run_indices, + runs, + run_offsets, + resume_by_scaffold.get(scaffold, {}), + scan_state, + ) + + region_result = fetchByRegion( + region, + WORKER_HEADER, + maf_fp, + entries, + WORKER_OUTPUT, + BATCHLOG, + as_fasta=WORKER_AS_FASTA, + fasta_header=WORKER_FASTA_HEADER, + expected_species=WORKER_EXPECTED_SPECIES, + fasta_dedupe=WORKER_FASTA_DEDUPE, + verbose=WORKER_VERBOSE, + warning_state=warning_state, + profile_state=profile_state, + scaffold_subdirs=WORKER_SCAFFOLD_SUBDIRS, + ) + + if scan_state and "resume" in scan_state: + # Resume the next region's scan from the first block that could + # still matter. If fetchByRegion returned before exhausting the + # iterator, "resume" is simply absent and the previous (earlier, + # conservative) offsets are kept -- correct, just more rescanning. + resume_by_scaffold[scaffold] = scan_state["resume"] + + batch_results[result_idx] = region_result + if not WORKER_SINGLE_OUTPUT: + # Field 4 is n.overlapping.blocks (see summary_headers). This + # used to read field 1, which is the region's START coordinate + # -- so every region beginning at coordinate 0 was counted as + # "no output" despite having been extracted fine, inflating + # zero_overlap_regions and tripping the zero-overlap failure + # threshold. A single-region BED starting at 0 always failed. + fields = region_result.split("\t") + if len(fields) >= 5: + try: + if int(fields[4]) == 0: + no_output_regions += 1 + else: + written_regions += 1 + except ValueError: + pass + finally: + index_fp.close() - region_result = fetchByRegion( - region, - WORKER_HEADER, - maf_fp, - WORKER_INDEX, - WORKER_OUTPUT, - BATCHLOG, - as_fasta=WORKER_AS_FASTA, - fasta_header=WORKER_FASTA_HEADER, - expected_species=WORKER_EXPECTED_SPECIES, - fasta_dedupe=WORKER_FASTA_DEDUPE, - verbose=WORKER_VERBOSE, - warning_state=warning_state, - profile_state=profile_state, - scaffold_subdirs=WORKER_SCAFFOLD_SUBDIRS, - ) - batch_results[result_idx] = region_result - if not WORKER_SINGLE_OUTPUT: - fields = region_result.split("\t") - if len(fields) >= 2: - try: - if int(fields[1]) == 0: - no_output_regions += 1 - else: - written_regions += 1 - except ValueError: - pass if profile_state is not None: profile_state["time_batch_total"] = time.perf_counter() - batch_timer_start @@ -1058,13 +1299,19 @@ def fetchByScaffold(scaffold, start_end, maf_file, maf_header, maf_compression, try: with COMMON.openMaf(maf_file, maf_compression, "rb") as mfp, open(output_path, "wb") as outfp: LOG.info(f">>> Scaffold {scaffold}"); - data = COMMON.readMafBlockBytes(mfp, maf_compression, start, end) outfp.write(maf_header.encode("utf-8")) - outfp.write(data) + # Streamed in bounded chunks: a scaffold can be hundreds of GB, + # far too large to materialize as one bytes object. + written = COMMON.copyMafRangeToStream(mfp, maf_compression, start, end, outfp) except Exception as e: - LOG.error(f"{scaffold}: {e}") + # repr(), plus the traceback, because some of the exceptions that + # actually occur here have an empty str() -- MemoryError being the + # one that made a real failure log as a blank message. + LOG.error(f"{scaffold}: FAILED to extract: {e!r}") + LOG.debug("Traceback for %s:\n%s", scaffold, traceback.format_exc()) + return {"scaffold": scaffold, "output_path": output_path, "ok": False, "error": repr(e)} - return f"{output_path}: Wrote {scaffold}"; + return {"scaffold": scaffold, "output_path": output_path, "ok": True, "bytes_written": written} def fetchScaffoldsSequential(ordered_scaffolds, index, maf_file, maf_header, maf_compression, out_dir, LOG, scaffold_subdirs=False): @@ -1081,13 +1328,16 @@ def fetchScaffoldsSequential(ordered_scaffolds, index, maf_file, maf_header, maf output_path = buildOutputPath(out_dir, scaffold, scaffold, ".maf", scaffold_subdirs) try: LOG.info(f">>> Scaffold {scaffold}"); - data = COMMON.readMafBlockBytes(mfp, maf_compression, start, end) with open(output_path, "wb") as outfp: outfp.write(maf_header.encode("utf-8")) - outfp.write(data) + # Streamed in bounded chunks -- see fetchByScaffold. + written = COMMON.copyMafRangeToStream(mfp, maf_compression, start, end, outfp) except Exception as e: - LOG.error(f"{scaffold}: {e}") - results.append(f"{output_path}: Wrote {scaffold}") + LOG.error(f"{scaffold}: FAILED to extract: {e!r}") + LOG.debug("Traceback for %s:\n%s", scaffold, traceback.format_exc()) + results.append({"scaffold": scaffold, "output_path": output_path, "ok": False, "error": repr(e)}) + continue + results.append({"scaffold": scaffold, "output_path": output_path, "ok": True, "bytes_written": written}) return results @@ -1157,7 +1407,55 @@ def run_fetch(args, cmdline="mafutils fetch"): LOG.info(f"Parsing index file.... {index_file}"); COMMON.validateIndexHeader(COMMON.readIndexHeader(index_file), args.maf_file, maf_compression, LOG, strict=args.verify_hash) - index = parseIndex(index_file, LOG, args.mode); + + # Scaffold mode's index is one line per scaffold (trivially small), so it is + # still loaded. Block mode streams its index instead -- see below. + index = parseIndex(index_file, LOG, args.mode) if args.mode == "scaffold" else None + + scaffold_runs = None + scaffold_first_run = {} + if args.mode == "block": + scaffold_index_file = COMMON.deriveScaffoldIndexPath(args.maf_file) + if not os.path.isfile(scaffold_index_file): + LOG.error( + f"Block mode needs the scaffold index ({scaffold_index_file}) to order regions " + f"the same way the block index is ordered, but it was not found. " + f"Run `mafutils index` to create it." + ) + sys.exit(1) + + # The two indexes must come from the same `mafutils index` run: block + # mode locates blocks using byte ranges recorded in the scaffold index, + # so a mismatched pair silently drops regions (an off-by-one boundary + # is enough to skip every run's first block). Both files stamp the + # MAF's size/mtime/hash in their header, so compare those -- the same + # cross-check `mafutils validate` performs. + block_header = COMMON.readIndexHeader(index_file) + scaffold_header = COMMON.readIndexHeader(scaffold_index_file) + if block_header is None or scaffold_header is None or block_header != scaffold_header: + LOG.error( + f"The block index ({index_file}) and scaffold index ({scaffold_index_file}) " + f"do not come from the same `mafutils index` run, so their byte offsets cannot " + f"be trusted together. Rebuild both with `mafutils index`." + ) + LOG.error(f" block index header: {block_header}") + LOG.error(f" scaffold index header: {scaffold_header}") + sys.exit(1) + + runs = readScaffoldRuns(scaffold_index_file) + # One binary search per run, once, to get a starting point for each + # run's records. These are scan hints, not exact bounds -- the run's + # MAF byte range is what actually delimits it. + run_offsets = computeRunIndexOffsets(index_file, runs) + scaffold_to_runs = {} + for i, (scaffold, _s, _e) in enumerate(runs): + scaffold_to_runs.setdefault(scaffold, []).append(i) + scaffold_first_run.setdefault(scaffold, i) + scaffold_runs = (scaffold_to_runs, runs, run_offsets) + LOG.info( + f"Loaded {len(runs)} scaffold run(s) covering {len(scaffold_to_runs)} scaffold(s) " + f"from {scaffold_index_file}" + ) LOG.info(f"Parsing BED file...... {args.bed_file}"); regions = parseBed(args.bed_file, LOG, args.mode); @@ -1208,6 +1506,7 @@ def run_fetch(args, cmdline="mafutils fetch"): # what lets a single reused handle avoid ever seeking backward. ordered_scaffolds = sorted(scaffold_set, key=lambda s: index[s][0]) + scaffold_results = [] if maf_compression == "gz": if args.processes > 1: LOG.warning( @@ -1215,8 +1514,7 @@ def run_fetch(args, cmdline="mafutils fetch"): "gzip-compressed files; --processes will be ignored. (Use a bgzip-compressed MAF " "for real parallel speedup on compressed input.)" ) - for result in fetchScaffoldsSequential(ordered_scaffolds, index, args.maf_file, maf_header, maf_compression, args.output, LOG, scaffold_subdirs=args.scaffold_subdirs): - LOG.info(result) + scaffold_results = fetchScaffoldsSequential(ordered_scaffolds, index, args.maf_file, maf_header, maf_compression, args.output, LOG, scaffold_subdirs=args.scaffold_subdirs) else: # Extract from MAF using the region index for matching scaffolds with ProcessPoolExecutor(max_workers=args.processes) as executor: @@ -1235,8 +1533,39 @@ def run_fetch(args, cmdline="mafutils fetch"): args.scaffold_subdirs, )) for future in futures: - result = future.result() - LOG.info(result) + scaffold_results.append(future.result()) + + # Only report "Wrote" for extractions that actually produced data -- + # never for one that raised or came back empty. + for result in scaffold_results: + if result["ok"] and result["bytes_written"] > 0: + LOG.info(f"{result['output_path']}: Wrote {result['scaffold']} ({result['bytes_written']} bytes)") + + # Any scaffold that failed to extract must fail the run. Previously + # these errors were logged and then reported as "Wrote " + # anyway, so fetch exited 0 with a 0-byte output file and downstream + # consumers proceeded on silently-corrupt state. + failed = [r for r in scaffold_results if not r["ok"]] + if failed: + LOG.error( + "Failed to extract %d of %d scaffold(s): %s", + len(failed), + len(scaffold_results), + ", ".join(r["scaffold"] for r in failed), + ) + sys.exit(1) + + empty = [r for r in scaffold_results if r["ok"] and r["bytes_written"] == 0] + if empty: + # No exception, but nothing came back for the requested range -- + # e.g. a bad/stale index offset. Still not a success. + LOG.error( + "Extracted 0 alignment bytes for %d scaffold(s): %s. Check that the index matches this MAF.", + len(empty), + ", ".join(r["scaffold"] for r in empty), + ) + sys.exit(1) + return; # Done with scaffold mode ############################## @@ -1261,18 +1590,35 @@ def run_fetch(args, cmdline="mafutils fetch"): "for real parallel speedup on compressed input.)" ) effective_processes = 1 - prefetch_cache = prefetchBlockCache(regions, index, args.maf_file, maf_compression, LOG) - # Every block any region needs is now decoded, in strictly-ascending - # file order, exactly once. initBatchWorker pre-seeds the worker's - # cache from this so fetchByRegion's lookups never seek at all. + # No prefetch pass any more. Regions are sorted into index order below + # and each worker walks the index forward, so blocks are read in + # ascending file order for EVERY compression type -- which is the + # property prefetchBlockCache existed to give gzip, now obtained + # structurally instead of by decoding every needed block into RAM + # first (a real run reported "Prefetching 1031475 distinct block(s)", + # i.e. its own O(needed blocks) memory blowup). Reuse of a block shared + # by neighbouring regions is still handled, by the bounded + # WORKER_BLOCK_CACHE LRU -- and sorting makes those regions adjacent, + # so a small cache captures it. + + # Sort regions into INDEX order (scaffold order as it appears in the MAF, + # then coordinate), so the block index can be streamed forward instead of + # loaded into memory. Regions are the cheap side of this: ~979k regions vs + # ~233M blocks on a whole-genome MAF. + unknown_rank = len(scaffold_first_run) + regions.sort(key=lambda r: (scaffold_first_run.get(r["scaffold"], unknown_rank), r["start"], r["end"])) num_regions = len(regions); - batch_size = pick_chunk_size(num_regions, effective_processes); - if batch_size >= num_regions: - batch_size = 1 + # One batch per worker: each streams its own span of the index. Splitting + # by region count (not by scaffold) keeps parallelism independent of how + # many scaffolds there are and of how wildly their sizes differ. + n_batches = max(1, min(effective_processes, num_regions)) + batch_size = (num_regions + n_batches - 1) // n_batches LOG.info(f"Using batch size of {batch_size} regions for {num_regions} total regions with {effective_processes} processes."); batches = list(chunker(regions, size=batch_size)) LOG.info(f"Divided {len(regions)} regions into {len(batches)} batches of up to {batch_size} regions each."); + + # Region batching for parallel processing if args.single_output: @@ -1327,6 +1673,8 @@ def run_fetch(args, cmdline="mafutils fetch"): zero_overlap_regions = 0 fail_fast_triggered = False fail_fast_message = None + no_overlap_threshold_warned = False + missing_scaffold_regions = 0 profile_totals = initProfileState() if args.profile else None profiled_batches = 0 with ProcessPoolExecutor( @@ -1336,7 +1684,7 @@ def run_fetch(args, cmdline="mafutils fetch"): maf_header, args.maf_file, maf_compression, - index, + None, # block mode streams the index; nothing to hand to workers args.output, args.single_output, args.fasta, @@ -1347,6 +1695,8 @@ def run_fetch(args, cmdline="mafutils fetch"): args.profile, prefetch_cache, args.scaffold_subdirs, + index_file, + scaffold_runs, ), ) as executor, open(info_outfile, "w") as info_out: summary_headers = ["scaffold", "start", "end", "basename", "n.overlapping.blocks", "block.lengths", "interblock.distances", "n.sequences"] @@ -1376,6 +1726,7 @@ def run_fetch(args, cmdline="mafutils fetch"): info_out.write(r + "\n") total_regions_reported += result["processed_regions"] zero_overlap_regions += result["zero_overlap_regions"] + missing_scaffold_regions += result["warning_counts"].get("missing-scaffold", 0) if result["warning_count"] > 0 and args.verbose: for warning_message in result["warning_messages"]: @@ -1396,19 +1747,27 @@ def run_fetch(args, cmdline="mafutils fetch"): formatWarningCounts(result["warning_counts"]), ) - if zero_overlap_regions > MAX_NO_OVERLAP_REGIONS_DEFAULT: - fail_fast_triggered = True - fail_fast_message = ( - f"Zero-overlap regions ({zero_overlap_regions}) exceeded hardcoded max " - f"({MAX_NO_OVERLAP_REGIONS_DEFAULT}). Failing run." - ) - elif zero_overlap_fraction_final_floor > MAX_NO_OVERLAP_FRACTION_DEFAULT: - fail_fast_triggered = True - fail_fast_message = ( - f"Zero-overlap fraction cannot recover below hardcoded max " - f"({MAX_NO_OVERLAP_FRACTION_DEFAULT:.6f}); current lower bound is " - f"{zero_overlap_fraction_final_floor:.6f}. Failing run." - ) + # Advisory only -- regions that overlap no alignment block + # are a normal outcome when fetching many intervals, so + # they warn rather than abort. (The genuinely diagnostic + # case, *every* region finding nothing, is still fatal -- + # see the end-of-run check.) Warn once, not per batch. + if not no_overlap_threshold_warned: + if zero_overlap_regions > ADVISORY_NO_OVERLAP_REGIONS: + no_overlap_threshold_warned = True + LOG.warning( + "Zero-overlap regions (%d) exceeded the advisory max (%d); continuing.", + zero_overlap_regions, + ADVISORY_NO_OVERLAP_REGIONS, + ) + elif zero_overlap_fraction_final_floor > ADVISORY_NO_OVERLAP_FRACTION: + no_overlap_threshold_warned = True + LOG.warning( + "Zero-overlap fraction is at least %.6f, above the advisory max " + "(%.6f); continuing.", + zero_overlap_fraction_final_floor, + ADVISORY_NO_OVERLAP_FRACTION, + ) else: results.append(result["results"]) @@ -1443,6 +1802,19 @@ def run_fetch(args, cmdline="mafutils fetch"): profile_totals["time_write_fasta"], ) + if not args.single_output and missing_scaffold_regions > 0: + # Advisory, like the other non-overlap outcomes: fetch is normally + # pointed at many intervals and a BED may legitimately carry scaffolds + # this MAF doesn't have. Still worth calling out prominently, since the + # usual cause is a naming mismatch (e.g. "chr1" vs "1") or the wrong + # index, and every region on that scaffold silently yields nothing. + LOG.warning( + "%d region(s) reference a scaffold that is not in the index at all. If that is " + "unexpected, check that the BED's scaffold names match the MAF's (and that this is " + "the right index).", + missing_scaffold_regions, + ) + if not args.single_output and total_regions_reported > 0 and zero_overlap_regions > 0: zero_overlap_fraction = zero_overlap_regions / total_regions_reported LOG.warning( @@ -1451,26 +1823,31 @@ def run_fetch(args, cmdline="mafutils fetch"): total_regions_reported, zero_overlap_fraction, ) + # Advisory too. A region overlapping no alignment block is normal data, + # not a malfunction -- fetch is typically pointed at many intervals, and + # a one-region BED that happens not to overlap is the same situation as + # "all regions" failing. Genuine misuse is caught earlier and loudly: + # index/MAF hash+size mismatch (validateIndexHeader), a block/scaffold + # index pair from different runs, and per-scaffold extraction failures. if zero_overlap_regions == total_regions_reported: - LOG.error("No regions overlapped any MAF block. Failing run.") - sys.exit(2) + LOG.warning("No regions overlapped any MAF block.") + # The count/fraction thresholds are advisory: some regions not + # overlapping any block is expected when fetching many intervals. if ( - MAX_NO_OVERLAP_REGIONS_DEFAULT >= 0 - and zero_overlap_regions > MAX_NO_OVERLAP_REGIONS_DEFAULT + ADVISORY_NO_OVERLAP_REGIONS >= 0 + and zero_overlap_regions > ADVISORY_NO_OVERLAP_REGIONS ): - LOG.error( - "Zero-overlap regions (%d) exceeded hardcoded max (%d). Failing run.", + LOG.warning( + "Zero-overlap regions (%d) exceeded the advisory max (%d).", zero_overlap_regions, - MAX_NO_OVERLAP_REGIONS_DEFAULT, + ADVISORY_NO_OVERLAP_REGIONS, ) - sys.exit(2) - if zero_overlap_fraction > MAX_NO_OVERLAP_FRACTION_DEFAULT: - LOG.error( - "Zero-overlap fraction (%.6f) exceeded hardcoded max (%.6f). Failing run.", + if zero_overlap_fraction > ADVISORY_NO_OVERLAP_FRACTION: + LOG.warning( + "Zero-overlap fraction (%.6f) exceeded the advisory max (%.6f).", zero_overlap_fraction, - MAX_NO_OVERLAP_FRACTION_DEFAULT, + ADVISORY_NO_OVERLAP_FRACTION, ) - sys.exit(2) # if args.single_output: # with open(output_file, "w", encoding="utf-8") as out_stream: diff --git a/mafutils/gc.py b/mafutils/gc.py index 43a4dfe..0a648c4 100644 --- a/mafutils/gc.py +++ b/mafutils/gc.py @@ -41,7 +41,8 @@ import os import sys from collections import Counter, defaultdict -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait +from itertools import islice from enum import Enum from types import SimpleNamespace from typing import Annotated, Optional @@ -105,12 +106,19 @@ def updateCountsFromSeqLine(line, counts): ############################################################################# -def parseIndex(index_file, LOG): +def iterIndexEntries(index_file, LOG): """ - Parses a block-level index file (same 8-column format as mafutils stats). - Only offset_start/offset_end are used here. + Streams (offset_start, offset_end) pairs from a block-level index file + (same 8-column format as mafutils stats). Only those two columns are + used here. + + A generator rather than a list on purpose: whole-genome indexes are far + too large to materialize. A real 241-species index had ~233M block + records, which as a list of tuples measured ~33GB of RAM (plus another + ~2GB once it was re-chunked by slicing) and OOM-killed gc at 62GB before + it computed anything. Streaming keeps memory flat regardless of genome + size. """ - entries = [] with open(index_file, "r", encoding="utf-8") as fp: for line in fp: if not line.strip() or line.startswith("#"): @@ -125,13 +133,22 @@ def parseIndex(index_file, LOG): except ValueError: LOG.warning(f"Skipping malformed index line: {line.strip()}") continue - entries.append((offset_start, offset_end)) - return entries + yield (offset_start, offset_end) -def chunker(seq, size): - for i in range(0, len(seq), size): - yield seq[i : i + size] +def chunker(iterable, size): + """ + Yields lists of up to `size` items from any iterable. Pulls lazily via + islice so it works on a generator without materializing the source -- + the previous list-slicing version required (and copied) the whole + sequence up front. + """ + iterator = iter(iterable) + while True: + chunk = list(islice(iterator, size)) + if not chunk: + return + yield chunk ############################################################################# @@ -165,40 +182,63 @@ def workerTask(task_id, maf_file, maf_compression, entries): return counts -def mergeCounts(results): - merged = defaultdict(newBaseCounts) - for counts in results: - for key, rec in counts.items(): - m = merged[key] - m["A"] += rec["A"] - m["C"] += rec["C"] - m["G"] += rec["G"] - m["T"] += rec["T"] - return merged +def mergeCountsInto(target, counts): + """ + Folds one worker's (species, chrom) base counts into an accumulator. + Used to merge each chunk's result as it arrives, rather than collecting + every worker's counts and merging at the end. + """ + for key, rec in counts.items(): + m = target[key] + m["A"] += rec["A"] + m["C"] += rec["C"] + m["G"] += rec["G"] + m["T"] += rec["T"] def runParallelGC(maf_file, maf_compression, index_file, processes, chunk_size, LOG, verify_hash=False): - LOG.info(f"Parsing index file: {index_file}") + LOG.info(f"Streaming index file: {index_file}") COMMON.validateIndexHeader(COMMON.readIndexHeader(index_file), maf_file, maf_compression, LOG, strict=verify_hash) - entries = parseIndex(index_file, LOG) - if not entries: - LOG.error("No valid index entries found.") - sys.exit(1) - LOG.info(f"Loaded {len(entries)} indexed blocks.") - chunks = list(chunker(entries, chunk_size)) - LOG.info(f"Running {len(chunks)} tasks with chunk size {chunk_size} across {processes} process(es).") + LOG.info(f"Streaming index in chunks of {chunk_size} blocks across {processes} process(es).") + + # Bound how many chunks are in flight at once. Submitting every chunk + # upfront would queue the whole index into the executor (and collecting + # every worker's result before merging would hold every partial count + # dict), which reintroduces the memory scaling this streaming is meant to + # avoid -- a whole-genome index is ~233M blocks / ~47k chunks. + max_in_flight = max(1, processes) * 2 + + merged = defaultdict(newBaseCounts) + total_blocks = 0 + tasks_done = 0 + chunk_iter = enumerate(chunker(iterIndexEntries(index_file, LOG), chunk_size), start=1) - results = [] with ProcessPoolExecutor(max_workers=processes) as executor: - futures = [ - executor.submit(workerTask, task_id, maf_file, maf_compression, chunk) - for task_id, chunk in enumerate(chunks, start=1) - ] - for future in futures: - results.append(future.result()) - - return mergeCounts(results) + in_flight = {} + for task_id, chunk in chunk_iter: + in_flight[executor.submit(workerTask, task_id, maf_file, maf_compression, chunk)] = len(chunk) + if len(in_flight) < max_in_flight: + continue + done, _pending = wait(list(in_flight.keys()), return_when=FIRST_COMPLETED) + for future in done: + # Merge as results arrive, then drop them, so peak memory is + # one accumulator plus the in-flight chunks. + total_blocks += in_flight.pop(future) + mergeCountsInto(merged, future.result()) + tasks_done += 1 + + for future in in_flight: + total_blocks += in_flight[future] + mergeCountsInto(merged, future.result()) + tasks_done += 1 + + if not total_blocks: + LOG.error("No valid index entries found.") + sys.exit(1) + + LOG.info(f"Processed {total_blocks} indexed blocks in {tasks_done} task(s).") + return merged ############################################################################# diff --git a/mafutils/lib/common.py b/mafutils/lib/common.py index 5742594..57fdabc 100644 --- a/mafutils/lib/common.py +++ b/mafutils/lib/common.py @@ -195,11 +195,62 @@ def computeFileHash(maf_file, algo="md5"): ############################################################################# +COPY_CHUNK_SIZE = 8 * 1024 * 1024 + + +def copyMafRangeToStream(handle, compression, offset_start, offset_end, out_stream, chunk_size=COPY_CHUNK_SIZE): + """ + Copies the byte range [offset_start, offset_end) from an open binary-mode + MAF handle straight into out_stream, in bounded chunks. Returns the number + of bytes written. + + Use this instead of readMafBlockBytes() whenever the range can be large -- + notably whole-scaffold extraction. readMafBlockBytes() materializes the + entire range as one bytes object, which is fine for a single alignment + block (KBs) but fatal for a scaffold: on a real 7.4TB 241-species MAF, + chr1 spans bytes 16 -> 597,495,640,164, so that became a single ~597GB + read() and died with a MemoryError whose str() is the empty string + (hence a completely blank error message in the log). + + Memory here is O(chunk_size) regardless of range size, for every + compression type -- including bgzip, whose line-by-line path must also + write as it goes rather than accumulating lines to join at the end. + """ + handle.seek(offset_start) + written = 0 + + if compression == "bgzip": + # Virtual offsets are comparable but not subtractable, so read + # forward until the handle's position reaches offset_end. Write each + # line straight out instead of accumulating them. + while handle.tell() < offset_end: + line = handle.readline() + if not line: + break + out_stream.write(line) + written += len(line) + return written + + remaining = offset_end - offset_start + while remaining > 0: + chunk = handle.read(min(chunk_size, remaining)) + if not chunk: + break + out_stream.write(chunk) + written += len(chunk) + remaining -= len(chunk) + return written + + def readMafBlockBytes(handle, compression, offset_start, offset_end): """ Seeks an open binary-mode MAF handle to offset_start and reads exactly the bytes of one block, returning them. + Only appropriate for genuinely block-sized ranges, since it holds the + whole range in memory -- see copyMafRangeToStream() for large ranges + such as whole scaffolds. + For "none"/"gz", offset_start/offset_end are raw byte positions / decompressed-stream positions respectively, so a simple subtraction gives the correct read length. diff --git a/mafutils/stats.py b/mafutils/stats.py index 8c855e9..fd68906 100644 --- a/mafutils/stats.py +++ b/mafutils/stats.py @@ -35,7 +35,8 @@ from types import SimpleNamespace from typing import Annotated, Optional from collections import defaultdict -from concurrent.futures import ProcessPoolExecutor +from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait +from itertools import islice import typer @@ -75,8 +76,24 @@ def parseExpectedSpecies(args): return species -def parseIndex(index_file, LOG): - entries = [] +def iterIndexEntries(index_file, LOG): + """ + Streams one tuple per block from the block index. + + A generator rather than a list on purpose: a whole-genome index has on + the order of 233M block records, and materializing them (as mafutils gc + used to do, with only 2 fields per block rather than the 8 here) measured + ~32GB of RAM and OOM-killed the run before any work happened. See + DEVELOPMENT.md. + + block_id numbering deliberately matches the previous list-building + version bit for bit, including its quirk: the counter is incremented + *before* the int() conversions, so a line that has enough fields but + fails to parse still consumes an id (leaving a gap in block_id) rather + than renumbering everything after it. Only matters for malformed + indexes, but block_id is written into the per-block output table, so + changing it would silently change that output. + """ with open(index_file, "r", encoding="utf-8") as fp: block_id = 0 for line in fp: @@ -88,26 +105,33 @@ def parseIndex(index_file, LOG): continue try: block_id += 1 - entries.append( - ( - block_id, - fields[0], # scaffold - int(fields[1]), # ref_start - int(fields[2]), # ref_length - int(fields[3]), # aln_length - int(fields[5]), # n_seq_lines (index) - int(fields[6]), # offset_start - int(fields[7]), # offset_end - ) + entry = ( + block_id, + fields[0], # scaffold + int(fields[1]), # ref_start + int(fields[2]), # ref_length + int(fields[3]), # aln_length + int(fields[5]), # n_seq_lines (index) + int(fields[6]), # offset_start + int(fields[7]), # offset_end ) except ValueError: LOG.warning(f"Skipping malformed index line: {line.strip()}") - return entries + continue + yield entry -def chunker(seq, size): - for i in range(0, len(seq), size): - yield seq[i : i + size] +def chunker(iterable, size): + """ + Yields lists of up to `size` items from any iterable, pulling lazily via + islice so a generator source is never materialized. + """ + iterator = iter(iterable) + while True: + chunk = list(islice(iterator, size)) + if not chunk: + return + yield chunk def speciesFromSrc(src): @@ -211,7 +235,19 @@ def workerTask( entries, write_block_rows, tmp_dir, + maf_fp=None, ): + """ + Processes one chunk of index entries. + + maf_fp lets a caller supply an already-open handle to reuse across + chunks. That matters for the single-process path: entries are in + ascending file order, so reusing one handle keeps every seek forward. + Reopening per chunk would restart gzip decompression from byte 0 each + time, turning the sequential gz path quadratic once the index is + streamed in many chunks rather than handed over as one giant chunk. + Parallel workers pass nothing and open their own handle, as before. + """ overall = { "total_blocks": 0, "total_alignment_columns": 0, @@ -249,7 +285,10 @@ def workerTask( block_tmp_path = os.path.join(tmp_dir, f"maf_stats.blocks.task{task_id}.tsv") block_tmp = open(block_tmp_path, "w", encoding="utf-8") - with COMMON.openMaf(maf_file, maf_compression, "rb") as maf_fp: + owns_handle = maf_fp is None + if owns_handle: + maf_fp = COMMON.openMaf(maf_file, maf_compression, "rb") + try: for entry in entries: ( block_id, @@ -322,11 +361,15 @@ def workerTask( ) + "\n" ) + finally: + if owns_handle: + maf_fp.close() if block_tmp is not None: block_tmp.close() return { + "task_id": task_id, "overall": overall, "species": dict(species), "observed_species": sorted(observed_species), @@ -334,41 +377,54 @@ def workerTask( } -def mergeStats(results): - overall = defaultdict(float) - species = defaultdict( - lambda: { - "blocks_present": 0, - "copy_lines": 0, - "duplicated_blocks": 0, - "duplicate_copies_total": 0, - "gaps_total": 0, - "nongaps_total": 0, - "gap_pct_block_sum": 0.0, - "all_gap_blocks": 0, - } - ) - observed_species = set() - block_tmp_paths = [] - - for res in results: - for k, v in res["overall"].items(): - if k == "site_gap_hist": - if "site_gap_hist" not in overall: - overall["site_gap_hist"] = [0] * (SITE_GAP_HIST_BINS + 1) - for i, count in enumerate(v): - overall["site_gap_hist"][i] += count - else: - overall[k] += v - for sp, rec in res["species"].items(): - s = species[sp] - for k, v in rec.items(): - s[k] += v - observed_species.update(res["observed_species"]) - if res["block_tmp_path"]: - block_tmp_paths.append(res["block_tmp_path"]) +def newStatsAccumulator(): + """ + Fresh accumulator for merging worker results one at a time. Paired with + mergeStatsInto() so each chunk's result can be folded in and dropped as + it arrives -- holding every chunk's result to merge at the end would + scale with index size (a whole-genome run is ~47k chunks, each carrying + a per-species dict). + """ + return { + "overall": defaultdict(float), + "species": defaultdict( + lambda: { + "blocks_present": 0, + "copy_lines": 0, + "duplicated_blocks": 0, + "duplicate_copies_total": 0, + "gaps_total": 0, + "nongaps_total": 0, + "gap_pct_block_sum": 0.0, + "all_gap_blocks": 0, + } + ), + "observed_species": set(), + "block_tmp_paths": [], + } + - return overall, species, observed_species, block_tmp_paths +def mergeStatsInto(acc, res): + overall = acc["overall"] + for k, v in res["overall"].items(): + if k == "site_gap_hist": + if "site_gap_hist" not in overall: + overall["site_gap_hist"] = [0] * (SITE_GAP_HIST_BINS + 1) + for i, count in enumerate(v): + overall["site_gap_hist"][i] += count + else: + overall[k] += v + for sp, rec in res["species"].items(): + s = acc["species"][sp] + for k, v in rec.items(): + s[k] += v + acc["observed_species"].update(res["observed_species"]) + if res["block_tmp_path"]: + # Keyed by task_id so the per-chunk block rows can be concatenated in + # index order. Results arrive out of order in the parallel path, and + # sorting by *filename* is wrong regardless: "task10" sorts before + # "task2", which scrambled block.tsv for any run with >=10 chunks. + acc["block_tmp_paths"].append((res["task_id"], res["block_tmp_path"])) def writeOverall(out_path, overall, n_species_total): @@ -478,7 +534,10 @@ def writeBlockTable(out_path, block_tmp_paths, n_species_total, expected_species with open(out_path, "w", encoding="utf-8") as out_fp: out_fp.write("\t".join(header) + "\n") - for tmp_path in sorted(block_tmp_paths): + # Already ordered by task id (= index order) by the caller; do NOT + # sort by path here -- filenames sort lexicographically, so "task10" + # would land before "task2". + for tmp_path in block_tmp_paths: with open(tmp_path, "r", encoding="utf-8") as in_fp: for line in in_fp: fields = line.rstrip("\n").split("\t") @@ -1188,13 +1247,8 @@ def run_stats(args, cmdline="mafutils stats"): if expected_species: LOG.info(f"Loaded {len(expected_species)} expected species for exact missing-species reporting.") - LOG.info(f"Parsing index file: {index_file}") + LOG.info(f"Reading index file: {index_file}") COMMON.validateIndexHeader(COMMON.readIndexHeader(index_file), args.maf_file, maf_compression, LOG, strict=args.verify_hash) - entries = parseIndex(index_file, LOG) - if not entries: - LOG.error("No valid index entries found.") - sys.exit(1) - LOG.info(f"Loaded {len(entries)} indexed blocks.") effective_processes = args.processes effective_chunk_size = args.chunk_size @@ -1206,22 +1260,32 @@ def run_stats(args, cmdline="mafutils stats"): "for real parallel speedup on compressed input.)" ) effective_processes = 1 - effective_chunk_size = len(entries) - # A single task covering every entry, in one worker: since entries are - # already in ascending file order, this keeps gzip seeking strictly - # forward (cheap) instead of paying a redundant-decompression cost - # per additional worker/chunk. + # gz previously forced a single chunk covering every entry, to keep + # one worker seeking strictly forward. That required knowing the + # total entry count upfront, which is exactly what streaming avoids. + # The single-process path below now reuses one open handle across + # chunks instead, which preserves forward-only seeking without + # materializing the index. - chunks = list(chunker(entries, effective_chunk_size)) - LOG.info(f"Running {len(chunks)} tasks with chunk size {effective_chunk_size} across {effective_processes} process(es).") + LOG.info(f"Streaming index in chunks of {effective_chunk_size} blocks across {effective_processes} process(es).") write_block_rows = (not args.no_block_table) or args.html_dashboard with tempfile.TemporaryDirectory(prefix="maf_stats_tmp_", dir=out_dir) as tmp_dir: - results = [] + acc = newStatsAccumulator() + total_blocks_seen = 0 + tasks_done = 0 + # The index is streamed, never materialized -- see iterIndexEntries. + chunk_iter = enumerate(chunker(iterIndexEntries(index_file, LOG), effective_chunk_size), start=1) + if effective_processes > 1: + # Bound in-flight work and merge results as they arrive, so peak + # memory is one accumulator plus the queued chunks rather than + # every chunk and every chunk's result. + max_in_flight = effective_processes * 2 with ProcessPoolExecutor(max_workers=effective_processes) as executor: - futures = [ - executor.submit( + in_flight = {} + for task_id, chunk in chunk_iter: + future = executor.submit( workerTask, task_id, args.maf_file, @@ -1230,10 +1294,17 @@ def run_stats(args, cmdline="mafutils stats"): write_block_rows, tmp_dir, ) - for task_id, chunk in enumerate(chunks, start=1) - ] - for f in futures: - results.append(f.result()) + in_flight[future] = len(chunk) + if len(in_flight) < max_in_flight: + continue + done, _pending = wait(list(in_flight.keys()), return_when=FIRST_COMPLETED) + for f in done: + total_blocks_seen += in_flight.pop(f) + mergeStatsInto(acc, f.result()) + tasks_done += 1 + for f in list(in_flight): + total_blocks_seen += in_flight.pop(f) + mergeStatsInto(acc, f.result()) else: # Skip ProcessPoolExecutor entirely at processes=1 -- pool # creation carries a large, mostly-fixed memory cost regardless @@ -1243,12 +1314,35 @@ def run_stats(args, cmdline="mafutils stats"): # in-process anyway. workerTask is a plain function with no # pool-initializer dependency, so calling it directly here is # equivalent to submitting it to a 1-worker pool. - for task_id, chunk in enumerate(chunks, start=1): - results.append( - workerTask(task_id, args.maf_file, maf_compression, chunk, write_block_rows, tmp_dir) - ) + # + # One handle is opened here and reused for every chunk: entries + # are in ascending file order, so all seeks stay forward. This is + # what preserves the cheap gz path now that the index arrives in + # many chunks instead of one -- reopening per chunk would restart + # gzip decompression from byte 0 each time. + with COMMON.openMaf(args.maf_file, maf_compression, "rb") as shared_fp: + for task_id, chunk in chunk_iter: + total_blocks_seen += len(chunk) + tasks_done += 1 + mergeStatsInto( + acc, + workerTask( + task_id, args.maf_file, maf_compression, chunk, + write_block_rows, tmp_dir, maf_fp=shared_fp, + ), + ) - overall, species, observed_species, block_tmp_paths = mergeStats(results) + if not total_blocks_seen: + LOG.error("No valid index entries found.") + sys.exit(1) + LOG.info(f"Processed {total_blocks_seen} indexed blocks in {tasks_done} tasks.") + + overall = acc["overall"] + species = acc["species"] + observed_species = acc["observed_species"] + # Sort numerically by task id so block rows land in index order + # regardless of which worker finished first. + block_tmp_paths = [path for _task_id, path in sorted(acc["block_tmp_paths"])] zero_seq_blocks = int(overall["blocks_with_zero_parsed_seq_lines_but_index_nonzero"]) if zero_seq_blocks: diff --git a/tests/example.maf.scaffold.idx b/tests/example.maf.scaffold.idx index abf1344..d9d3c1d 100644 --- a/tests/example.maf.scaffold.idx +++ b/tests/example.maf.scaffold.idx @@ -1,6 +1,7 @@ -chr1 16 264 -chr2 264 339 -chr3 339 483 -chr4 483 629 -chrX 629 701 -chr4 701 777 +# mafutils-index format=2 maf=example.maf compression=none size=776 mtime=1757533321.0 hash=md5:0c80ee622a5f9e5ef29d331e18278884 +chr1 16 263 +chr2 263 338 +chr3 338 482 +chr4 482 628 +chrX 628 700 +chr4 700 776 diff --git a/tests/test_compression.py b/tests/test_compression.py index 1e2a151..dd2b9f6 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -156,7 +156,6 @@ def test_fetch_gzip_falls_back_and_matches(tmp_path): assert result.returncode == 0, result.stderr output = combined_output(result) assert "inefficient on" in output - assert "Prefetching" in output for name in os.listdir(baseline_dir): if not name.endswith(".maf"): diff --git a/tests/test_fetch.py b/tests/test_fetch.py index 8000cd1..6e5c627 100644 --- a/tests/test_fetch.py +++ b/tests/test_fetch.py @@ -12,13 +12,18 @@ EXPECTED_MAF_DIR = os.path.join(TEST_DIR, "expected-maf") EXPECTED_FASTA_DIR = os.path.join(TEST_DIR, "expected-fasta") -# Regions with no overlapping alignment block are a normal, expected outcome -# (fetch typically processes many regions per run) -- fetch warns and moves -# on rather than failing. This is distinct from 09-missingchrom, where the -# BED references a scaffold that isn't in the index at all, which is a -# harder error and does exit non-zero (see the empty-expected-file branch -# below). -WARN_ONLY_NO_OVERLAP_REGIONS = {"06-noalign"} +# Regions that produce no output are a normal, expected outcome -- fetch is +# typically pointed at many intervals, so it warns and moves on rather than +# failing. Maps region id -> the warning code fetch should report: +# 06-noalign scaffold is in the index, but no block overlaps the region +# 09-missingchrom the BED names a scaffold the index doesn't have at all +# Both warn and exit 0. (09-missingchrom used to exit non-zero, but only as a +# side effect of the no-overlap count thresholds, which were being computed +# from the wrong summary column; those are advisory now.) +WARN_ONLY_REGIONS = { + "06-noalign": "no-overlap", + "09-missingchrom": "missing-scaffold", +} def get_regions(): """Yields (chrom, start, end, id) for each region in the BED file.""" @@ -51,17 +56,17 @@ def test_maf_fetch_region(chrom, start, end, region_id, tmp_path): expected_file = os.path.join(EXPECTED_MAF_DIR, f"{region_id}.maf") result = subprocess.run(cmd, check=False, cwd=REPO_ROOT, capture_output=True, text=True) - if region_id in WARN_ONLY_NO_OVERLAP_REGIONS: + if region_id in WARN_ONLY_REGIONS: + expected_warning = WARN_ONLY_REGIONS[region_id] assert result.returncode == 0, f"Expected exit 0 (warn, not fail) for region {region_id}: {result.stderr}" assert not os.path.exists(output_file), f"Did not expect output file for region {region_id}" - assert "no-overlap" in result.stderr, f"Expected a no-overlap warning for region {region_id}" - return - - if os.path.getsize(expected_file) == 0: - assert result.returncode != 0, f"Expected non-zero exit for region {region_id}" - assert not os.path.exists(output_file), f"Did not expect output file for region {region_id}" + assert expected_warning in result.stderr, f"Expected a {expected_warning} warning for region {region_id}" return + assert os.path.getsize(expected_file) > 0, ( + f"{expected_file} is empty, but every remaining region is expected to produce output; " + f"regions that legitimately produce none belong in WARN_ONLY_REGIONS" + ) assert result.returncode == 0, f"Unexpected non-zero exit for region {region_id}: {result.returncode}" got = open(output_file).read() @@ -97,17 +102,17 @@ def test_maf_fetch_fasta_region(chrom, start, end, region_id, tmp_path): expected_file = os.path.join(EXPECTED_FASTA_DIR, f"{region_id}.fa") result = subprocess.run(cmd, check=False, cwd=REPO_ROOT, capture_output=True, text=True) - if region_id in WARN_ONLY_NO_OVERLAP_REGIONS: + if region_id in WARN_ONLY_REGIONS: + expected_warning = WARN_ONLY_REGIONS[region_id] assert result.returncode == 0, f"Expected exit 0 (warn, not fail) for region {region_id}: {result.stderr}" assert not os.path.exists(output_file), f"Did not expect FASTA output file for region {region_id}" - assert "no-overlap" in result.stderr, f"Expected a no-overlap warning for region {region_id}" - return - - if os.path.getsize(expected_file) == 0: - assert result.returncode != 0, f"Expected non-zero FASTA exit for region {region_id}" - assert not os.path.exists(output_file), f"Did not expect FASTA output file for region {region_id}" + assert expected_warning in result.stderr, f"Expected a {expected_warning} warning for region {region_id}" return + assert os.path.getsize(expected_file) > 0, ( + f"{expected_file} is empty, but every remaining region is expected to produce output; " + f"regions that legitimately produce none belong in WARN_ONLY_REGIONS" + ) assert result.returncode == 0, f"Unexpected non-zero FASTA exit for region {region_id}: {result.returncode}" got = open(output_file).read() @@ -143,9 +148,9 @@ def test_scaffold_subdirs_groups_output_by_scaffold(tmp_path): result = subprocess.run(cmd, check=False, cwd=REPO_ROOT, capture_output=True, text=True) assert result.returncode == 0, f"Unexpected non-zero exit: {result.stderr}" - # chr1 has 3 regions with real output (06-noalign has no overlap, so no + # chr1 has 3 regions with real output (06-noalign overlaps nothing, so no # file); chr4 has 2; chr2/chr3/chrX have 1 each. chrZ (09-missingchrom) - # never gets that far (unknown scaffold errors before any writes). + # isn't in the index, so it produces no file and no chrZ/ subdirectory. assert sorted(os.listdir(os.path.join(out_dir, "chr1"))) == [ "01-singleblock-gap.maf", "02-truncatedblock.maf", diff --git a/tests/test_fetch_streaming.py b/tests/test_fetch_streaming.py new file mode 100644 index 0000000..8eefe6a --- /dev/null +++ b/tests/test_fetch_streaming.py @@ -0,0 +1,177 @@ +import os +import subprocess +import sys + +import pytest + +TEST_DIR = os.path.dirname(__file__) +REPO_ROOT = os.path.abspath(os.path.join(TEST_DIR, "..")) + +MAF_FILE = os.path.join(TEST_DIR, "example.maf") +INDEX_FILE = os.path.join(TEST_DIR, "example.maf.block.idx") +SCAFFOLD_INDEX = os.path.join(TEST_DIR, "example.maf.scaffold.idx") +BED_FILE = os.path.join(TEST_DIR, "example.bed") + +# 8 real blocks on one scaffold (CM000994.3), extracted from production data. +REAL_MAF = os.path.join(TEST_DIR, "real-excerpt.maf") +REAL_INDEX = os.path.join(TEST_DIR, "real-excerpt.maf.block.idx") + + +def run_fetch(args): + return subprocess.run( + [sys.executable, "-m", "mafutils", "fetch"] + args, + check=False, cwd=REPO_ROOT, capture_output=True, text=True, + ) + + +def outputs(directory, ext=".maf"): + if not os.path.isdir(directory): + return {} + return { + name: open(os.path.join(directory, name), "rb").read() + for name in os.listdir(directory) + if name.endswith(ext) + } + + +# --------------------------------------------------------------------------- +# Block mode streams the index, so results must not depend on how regions are +# partitioned across workers, nor on the order they appear in the BED. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("processes", [1, 2, 3, 8]) +def test_output_is_invariant_to_process_count(processes, tmp_path): + baseline = os.path.join(str(tmp_path), "p1") + assert run_fetch([MAF_FILE, BED_FILE, "--index", INDEX_FILE, "-b", "id", "-o", baseline, "-p", "1"]).returncode == 0 + + out = os.path.join(str(tmp_path), f"p{processes}") + result = run_fetch([MAF_FILE, BED_FILE, "--index", INDEX_FILE, "-b", "id", "-o", out, "-p", str(processes)]) + assert result.returncode == 0, result.stderr + assert outputs(out) == outputs(baseline) + + +def test_output_is_invariant_to_bed_order(tmp_path): + """ + Regions are sorted into index order internally, so a shuffled BED must + produce byte-identical per-region output. (Guards the sort/merge-join: if + the streaming assumed BED order, a reversed BED would seek backwards and + silently drop blocks.) + """ + with open(BED_FILE) as fp: + lines = [ln for ln in fp if ln.strip()] + + forward = os.path.join(str(tmp_path), "forward") + assert run_fetch([MAF_FILE, BED_FILE, "--index", INDEX_FILE, "-b", "id", "-o", forward]).returncode == 0 + + reversed_bed = os.path.join(str(tmp_path), "reversed.bed") + with open(reversed_bed, "w") as fp: + # Normalize line endings: example.bed's last line has no trailing + # newline, so writing the reversed list verbatim would splice two + # records together. + fp.write("".join(ln.rstrip("\n") + "\n" for ln in reversed(lines))) + + out = os.path.join(str(tmp_path), "reversed") + result = run_fetch([MAF_FILE, reversed_bed, "--index", INDEX_FILE, "-b", "id", "-o", out]) + assert result.returncode == 0, result.stderr + assert outputs(out) == outputs(forward) + + +def test_scaffold_appearing_in_multiple_runs_is_fully_covered(tmp_path): + """ + example.maf interleaves scaffolds (chr4 ... chrX ... chr4), so chr4 occupies + two separate runs of the index. A streaming implementation that assumes one + contiguous run per scaffold silently loses the later blocks -- 08-span- + multiple-gaps lives in chr4's SECOND run, so its output is the canary. + """ + with open(SCAFFOLD_INDEX) as fp: + scaffolds = [ln.split("\t")[0] for ln in fp if not ln.startswith("#") and ln.strip()] + assert len(scaffolds) > len(set(scaffolds)), "fixture no longer exercises multi-run scaffolds" + + bed = os.path.join(str(tmp_path), "chr4.bed") + with open(bed, "w") as fp: + fp.write("chr4\t4\t8\tfirst-run\n") # chr4's first run + fp.write("chr4\t31\t33\tsecond-run\n") # chr4's second run, after chrX + + out = os.path.join(str(tmp_path), "out") + result = run_fetch([MAF_FILE, bed, "--index", INDEX_FILE, "-b", "id", "-o", out]) + assert result.returncode == 0, result.stderr + + for name in ("first-run.maf", "second-run.maf"): + path = os.path.join(out, name) + assert os.path.isfile(path) and os.path.getsize(path) > 0, f"{name} missing/empty" + assert "chr4" in open(path).read() + + +# --------------------------------------------------------------------------- +# Single large scaffold: the degenerate shape for a streaming design. +# --------------------------------------------------------------------------- + + +def test_many_regions_on_one_scaffold_use_all_workers(tmp_path): + """All 8 blocks of real-excerpt.maf are on one scaffold; splitting by region + count (not by scaffold) must still spread them across workers.""" + bed = os.path.join(str(tmp_path), "many.bed") + with open(bed, "w") as fp: + for i in range(8): + start = 125888838 + i * 100 + fp.write(f"CM000994.3\t{start}\t{start + 50}\tr{i}\n") + + single = os.path.join(str(tmp_path), "p1") + assert run_fetch([REAL_MAF, bed, "--index", REAL_INDEX, "-b", "id", "-o", single, "-p", "1"]).returncode == 0 + + out = os.path.join(str(tmp_path), "p4") + result = run_fetch([REAL_MAF, bed, "--index", REAL_INDEX, "-b", "id", "-o", out, "-p", "4"]) + assert result.returncode == 0, result.stderr + assert "into 4 batches" in (result.stdout + result.stderr) + assert outputs(out) == outputs(single) + + +def test_single_region_spanning_whole_scaffold(tmp_path): + """ + One BED region covering an entire scaffold. The blocks for it are streamed + rather than gathered into a list, so this must work without holding every + overlapping block in memory. + """ + bed = os.path.join(str(tmp_path), "whole.bed") + with open(bed, "w") as fp: + fp.write("CM000994.3\t0\t250000000\twhole\n") + + out = os.path.join(str(tmp_path), "out") + result = run_fetch([REAL_MAF, bed, "--index", REAL_INDEX, "-b", "id", "-o", out]) + assert result.returncode == 0, result.stderr + + content = open(os.path.join(out, "whole.maf")).read() + # All 8 blocks of the excerpt fall inside this range. + assert content.count("\na") + content.startswith("a") == 8, "expected every block in the scaffold" + + +# --------------------------------------------------------------------------- +# The block/scaffold index pair must be consistent, and say so if not. +# --------------------------------------------------------------------------- + + +def test_mismatched_index_pair_fails_loudly(tmp_path): + """ + Block mode locates blocks via byte ranges recorded in the scaffold index, so + a pair from different `mafutils index` runs cannot be trusted together. An + off-by-one boundary alone is enough to skip every run's first block, which + would silently drop regions -- so it must be an error, not a warning. + """ + maf = os.path.join(str(tmp_path), "t.maf") + open(maf, "wb").write(open(MAF_FILE, "rb").read()) + open(maf + ".block.idx", "wb").write(open(INDEX_FILE, "rb").read()) + + # Scaffold index whose header disagrees with the block index's. + with open(SCAFFOLD_INDEX) as src, open(maf + ".scaffold.idx", "w") as dst: + for line in src: + dst.write(line.replace("hash=md5:", "hash=md5:deadbeef") if line.startswith("#") else line) + + bed = os.path.join(str(tmp_path), "r.bed") + with open(bed, "w") as fp: + fp.write("chr1\t2\t6\tr1\n") + + result = run_fetch([maf, bed, "--index", maf + ".block.idx", "-b", "id", + "-o", os.path.join(str(tmp_path), "out")]) + assert result.returncode != 0 + assert "do not come from the same" in (result.stdout + result.stderr) diff --git a/tests/test_scaffold_mode.py b/tests/test_scaffold_mode.py new file mode 100644 index 0000000..e77c5b4 --- /dev/null +++ b/tests/test_scaffold_mode.py @@ -0,0 +1,116 @@ +import io +import os +import subprocess +import sys + +import pytest + +from mafutils.lib import common as COMMON + +TEST_DIR = os.path.dirname(__file__) +REPO_ROOT = os.path.abspath(os.path.join(TEST_DIR, "..")) + +MAF_FILE = os.path.join(TEST_DIR, "example.maf") +SCAFFOLD_INDEX = os.path.join(TEST_DIR, "example.maf.scaffold.idx") + +HEADER = "# mafutils-index format=2 maf=t.maf compression=none size=776 mtime=1.0 hash=md5:x\n" + + +def run_fetch(args): + return subprocess.run( + [sys.executable, "-m", "mafutils", "fetch"] + args, + check=False, cwd=REPO_ROOT, capture_output=True, text=True, + ) + + +def write_bed(path, lines): + with open(path, "w") as fp: + for line in lines: + fp.write(line + "\n") + + +# --------------------------------------------------------------------------- +# copyMafRangeToStream: the streaming replacement for a single giant read +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("compression,fixture", [ + ("none", "example.maf"), + ("gz", "example.maf.gz"), + ("bgzip", "example.maf.bgz"), +]) +@pytest.mark.parametrize("chunk_size", [1, 7, 64, 8 * 1024 * 1024]) +def test_copy_range_matches_read_bytes_at_any_chunk_size(compression, fixture, chunk_size): + """ + Streaming a range in bounded chunks must produce exactly what the + single-read version produces, including when the chunk size is far + smaller than the range (so the copy loop runs many times) -- the real + fixtures are small enough that a default chunk size would only ever + loop once, which would leave the multi-chunk path untested. + """ + maf = os.path.join(TEST_DIR, fixture) + index = maf + ".block.idx" + + with open(index) as fp: + rows = [ln.split("\t") for ln in fp if ln.strip() and not ln.startswith("#")] + start, end = int(rows[0][6]), int(rows[-1][7]) # span many blocks at once + + with COMMON.openMaf(maf, compression, "rb") as handle: + expected = COMMON.readMafBlockBytes(handle, compression, start, end) + + buf = io.BytesIO() + with COMMON.openMaf(maf, compression, "rb") as handle: + written = COMMON.copyMafRangeToStream(handle, compression, start, end, buf, chunk_size=chunk_size) + + assert buf.getvalue() == expected + assert written == len(expected) + + +# --------------------------------------------------------------------------- +# fetch -m scaffold: must fail loudly, never report success on failure +# --------------------------------------------------------------------------- + + +def test_scaffold_mode_succeeds_normally(tmp_path): + bed = os.path.join(str(tmp_path), "s.bed") + write_bed(bed, ["chr1\t0\t50\tchr1", "chr4\t0\t50\tchr4"]) + out = os.path.join(str(tmp_path), "out") + + result = run_fetch([MAF_FILE, bed, "-i", SCAFFOLD_INDEX, "-m", "scaffold", "-o", out]) + assert result.returncode == 0, result.stderr + + for scaffold in ("chr1", "chr4"): + path = os.path.join(out, f"{scaffold}.maf") + assert os.path.getsize(path) > 0, f"{scaffold} output is empty" + with open(path) as fp: + assert fp.read().lstrip().startswith("#") or "a" in fp.read(), path + + +def test_scaffold_mode_exits_nonzero_when_extraction_yields_nothing(tmp_path): + """ + Regression test for a silent-failure bug: a scaffold that produced no + output was logged as an error but still reported as "Wrote ", + and fetch exited 0 -- so a downstream pipeline consumed a 0-byte file + believing it had succeeded. Simulated here with a stale index whose + offsets point past EOF; the real trigger was a MemoryError from trying + to read a 597GB scaffold in one allocation. + """ + maf = os.path.join(str(tmp_path), "t.maf") + with open(MAF_FILE) as src, open(maf, "w") as dst: + dst.write(src.read()) + + stale_index = os.path.join(str(tmp_path), "stale.scaffold.idx") + with open(stale_index, "w") as fp: + fp.write(HEADER) + fp.write("chr1\t900000\t999999\n") + + bed = os.path.join(str(tmp_path), "s.bed") + write_bed(bed, ["chr1\t0\t50\tchr1"]) + + result = run_fetch([maf, bed, "-i", stale_index, "-m", "scaffold", + "-o", os.path.join(str(tmp_path), "out")]) + combined = result.stdout + result.stderr + + assert result.returncode != 0, "a scaffold that extracted nothing must fail the run" + assert "Wrote chr1" not in combined, "must not claim it wrote a scaffold it failed to extract" + assert "0 alignment bytes" in combined diff --git a/tests/test_stats.py b/tests/test_stats.py index 4d8bbc1..2ce2569 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -126,6 +126,29 @@ def test_sequential_and_parallel_paths_produce_identical_output(tmp_path): assert read_species(seq_prefix + ".species.tsv") == read_species(par_prefix + ".species.tsv") +@pytest.mark.parametrize("processes,chunk_size", [(1, 1), (1, 3), (2, 1), (2, 2), (4, 3), (1, 10000)]) +def test_output_is_invariant_to_chunk_size_and_process_count(processes, chunk_size, tmp_path): + """ + The index is streamed in chunks and each chunk's result is merged as it + arrives, so results must not depend on how the stream happens to be + divided. Tiny chunk sizes force many chunks (and, in the parallel path, + exercise the bounded in-flight/merge-on-arrival loop) where the default + would produce a single chunk and test none of it. + """ + baseline_prefix = os.path.join(str(tmp_path), "baseline") + run_stats([MAF_FILE, INDEX_FILE, "-o", baseline_prefix, "-p", "1", "--chunk-size", "10000"]) + + out_prefix = os.path.join(str(tmp_path), f"p{processes}c{chunk_size}") + run_stats([MAF_FILE, INDEX_FILE, "-o", out_prefix, "-p", str(processes), "--chunk-size", str(chunk_size)]) + + assert read_overall(out_prefix + ".overall.tsv") == read_overall(baseline_prefix + ".overall.tsv") + assert read_species(out_prefix + ".species.tsv") == read_species(baseline_prefix + ".species.tsv") + with open(out_prefix + ".block.tsv") as a, open(baseline_prefix + ".block.tsv") as b: + # Block rows are written per chunk into temp files and concatenated, + # so this also guards their ordering/block_id numbering. + assert a.read() == b.read() + + def test_expected_species_file_reports_missing_species(tmp_path): species_file = os.path.join(str(tmp_path), "species.txt") with open(species_file, "w", encoding="utf-8") as fp: diff --git a/tests/test_validate.py b/tests/test_validate.py index eafd528..e8ca9a6 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -14,10 +14,25 @@ MAF_FILE = os.path.join(TEST_DIR, "example.maf") INDEX_FILE = os.path.join(TEST_DIR, "example.maf.block.idx") -# Deliberately preserved pre-hash-header fixture (see DEVELOPMENT.md) -- used -# directly as a "headerless for these fields" case rather than constructing -# one by hand. -HEADERLESS_SCAFFOLD_INDEX = os.path.join(TEST_DIR, "example.maf.scaffold.idx") + +def writeHeaderlessIndex(src_index, dst_path): + """ + Copies an index without its `# mafutils-index ...` header line, i.e. what + an index built by a pre-header mafutils version looks like. + + Built here rather than relying on a committed stale fixture: that made the + file's out-of-date-ness load-bearing but invisible, so regenerating it (as + the block-index streaming work required) silently changed what these tests + exercised. Mirrors how + test_verify_hash_errors_when_header_present_but_no_hash_field constructs + its own sibling case. + """ + with open(src_index) as src, open(dst_path, "w") as dst: + for line in src: + if line.startswith("# mafutils-index"): + continue + dst.write(line) + return dst_path def run_module(args, cwd=REPO_ROOT): @@ -58,12 +73,9 @@ def test_hash_is_reproducible(tmp_path): def test_validate_clean_match_exits_0(tmp_path): - # A freshly-built block+scaffold pair, unlike the checked-in MAF_FILE - # fixture whose default-path scaffold index is deliberately kept - # headerless (see HEADERLESS_SCAFFOLD_INDEX above) to exercise that path - # elsewhere -- that fixture would (correctly) make this case - # UNVERIFIABLE rather than VERIFIED, since it can't cross-check a - # headerless scaffold index. + # Build a fresh block+scaffold pair in tmp_path rather than reusing the + # checked-in fixtures, so this asserts on a pair that is definitionally + # consistent with each other and with the MAF. maf_copy = os.path.join(str(tmp_path), "clean.maf") shutil.copyfile(MAF_FILE, maf_copy) index_result = run_module(["index", maf_copy]) @@ -142,7 +154,8 @@ def test_validate_content_change_exits_1(tmp_path): def test_validate_headerless_index_exits_2(tmp_path): - result = run_module(["validate", MAF_FILE, HEADERLESS_SCAFFOLD_INDEX]) + headerless = writeHeaderlessIndex(INDEX_FILE, os.path.join(str(tmp_path), "headerless.block.idx")) + result = run_module(["validate", MAF_FILE, headerless]) assert result.returncode == 2, combined_output(result) assert "UNVERIFIABLE" in combined_output(result) @@ -178,9 +191,10 @@ def test_fetch_verify_hash_errors_on_fully_headerless_index(tmp_path): with open(bed_path, "w") as fp: fp.write("chr1\t0\t10\tregion1\n") + headerless = writeHeaderlessIndex(INDEX_FILE, os.path.join(str(tmp_path), "headerless.block.idx")) result = run_module([ "fetch", MAF_FILE, bed_path, - "--index", HEADERLESS_SCAFFOLD_INDEX, + "--index", headerless, "--verify-hash", "-o", os.path.join(str(tmp_path), "out"), ])