diff --git a/docs/advanced/ploidy-and-sex-chroms.md b/docs/advanced/ploidy-and-sex-chroms.md index 987b991..71108e1 100644 --- a/docs/advanced/ploidy-and-sex-chroms.md +++ b/docs/advanced/ploidy-and-sex-chroms.md @@ -19,6 +19,8 @@ AFQuery computes ploidy-aware AN for sex chromosomes (chrX, chrY) and the mitoch For each eligible sample at a given position, AFQuery adds the appropriate ploidy count to AN based on the sample's sex and the chromosome/position. +A sample contributing ploidy 0 is not eligible at that position at all: a female has no chrY to genotype, so she is neither a carrier nor homozygous reference there. On chrY, `n_eligible`, `N_HOM_REF` and `AN` therefore count males only. + --- ## Pseudoautosomal Regions (PAR) diff --git a/docs/guides/update-database.md b/docs/guides/update-database.md index b0b2717..35606f0 100644 --- a/docs/guides/update-database.md +++ b/docs/guides/update-database.md @@ -38,6 +38,8 @@ afquery update-db \ The new manifest follows the same format as the original (see [Manifest Format](manifest-format.md)). New samples are assigned monotonically increasing sample IDs. +New variants are merged into the per-chromosome bucket files the database already uses; buckets are created on demand, and no new top-level Parquet file is produced. See [Data Model](../reference/data-model.md#storage-layouts). + To add multiple manifests at once: ```bash @@ -67,9 +69,15 @@ decisions are comparable across batches. When new carriers push a partially-covered tech above the `--min-covered` threshold at positions that were previously below it, those positions are re-evaluated and their non-carrier samples once again count as `N_HOM_REF` -instead of `N_NO_COVERAGE`. The recomputation runs only for chromosomes -touched by the new samples; existing rows on other chromosomes are not -rewritten. +instead of `N_NO_COVERAGE`. Because that value derives from the tech bitmaps of +the whole cohort rather than from which file received rows, the recomputation +spans **every bucket of every chromosome the database holds**, not only the ones +the new samples carry variants on. Skipping the rest would leave an added WES +sample counted as `N_HOM_REF` everywhere else its capture BED reaches. Only a +batch that puts a sample into a capture-based technology can move those bitmaps, +so a WGS-only batch still visits nothing beyond its own chromosomes; adding a +WES sample to a large database rewrites broadly and takes minutes rather than +seconds. VCFs added via `update-db` should preserve `FORMAT/DP` and `FORMAT/GQ` (the bundled `resources/normalize_vcf.sh` does so by default). Samples without diff --git a/docs/reference/data-model.md b/docs/reference/data-model.md index de1ce99..d726731 100644 --- a/docs/reference/data-model.md +++ b/docs/reference/data-model.md @@ -22,6 +22,18 @@ This page documents the on-disk layout of an AFQuery database, including file fo └── wes_v2.pkl ``` +### Storage layouts + +`create-db` always writes the bucketed layout shown above. A single-file-per-chromosome +layout (`variants/chr1.parquet`, with no bucket directory) is also readable, and appears +in small hand-built databases and in test fixtures. `update-db` merges into whichever +layout a chromosome already uses, and creates buckets for a chromosome new to the +database. + +A chromosome must never have both. `afquery check` reports that as an error, because +queries resolve the bucket directory first and would silently ignore the flat file — and +every sample stored only in it. + --- ## manifest.json @@ -156,6 +168,10 @@ Variants are partitioned into 1-Mbp buckets: bucket_id = pos // 1_000_000 ``` +`update-db --add-samples` writes new positions into the bucket that owns +them, creating `bucket_N.parquet` when a batch extends a chromosome past its +previous last bucket. + !!! warning "DuckDB integer arithmetic" When computing bucket IDs in DuckDB SQL, always use the integer-division operator: ```sql diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 6eee06f..7ff8ca8 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -81,6 +81,84 @@ AFQuery requires DuckDB to use Parquet for all temporary files. Arrow IPC is not --- +## Samples Added by update-db Are Missing From Queries + +**Symptom:** After `afquery update-db --add-samples`, the new samples show up in +`afquery info --db ./db/ --samples` and `AN` grows by the expected amount, but they +never appear as carriers: `AC` does not increase at variants you know they carry, and +`afquery variant-info` does not list them. Allele frequencies across the whole database +drift downward after every update. + +`afquery check --db ./db/` reports one error per affected chromosome: + +``` +ERROR chr1: both variants/chr1/ (bucketed) and variants/chr1.parquet exist. +``` + +**Cause:** `update-db --add-samples` before 0.4.1 only understood the +single-file-per-chromosome variant layout. `create-db` produces the bucketed layout +(`variants//bucket_N.parquet`), so the merge found nothing to merge and wrote the +new samples to a fresh `variants/.parquet` instead. Queries read the bucketed +files and ignore that one, so the added samples counted toward `AN` through the capture +index but never as carriers — they were treated as homozygous reference everywhere. + +**Fix:** Upgrade to 0.4.1 or later, then repair the database. `afquery info --db ./db/ +--changelog` lists the samples added by each `add_samples` event. + +```bash +# 0. Stop all writers and back up the small files. +cp ./db/manifest.json ./db/manifest.json.bak +cp ./db/metadata.sqlite ./db/metadata.sqlite.bak + +# 1. Record the affected samples BEFORE removing them: removal deletes their +# phenotype rows, and you need them to rebuild the manifest. +sqlite3 ./db/metadata.sqlite \ + "SELECT s.sample_name, s.sex, t.tech_name, s.vcf_path, + group_concat(p.phenotype_code) + FROM samples s + JOIN technologies t ON s.tech_id = t.tech_id + LEFT JOIN sample_phenotype p ON p.sample_id = s.sample_id + WHERE s.sample_name IN ('SAMPLE_1','SAMPLE_2') GROUP BY s.sample_id;" + +# 2. List the flat files before deleting anything. Only those with a sibling +# directory of the same name are orphans; one without a sibling is the only +# copy of that chromosome and queries still read it. +for f in ./db/variants/*.parquet; do + [ -d "${f%.parquet}" ] && echo "orphan: $f" || echo "KEEP (no sibling): $f" +done + +# 3. Remove the affected samples. This clears their bits from both layouts and is +# safe on a split database. Every bucket is read, so budget minutes, not +# seconds, and do not interrupt it. +afquery update-db --db ./db/ --remove-samples SAMPLE_1 --remove-samples SAMPLE_2 + +# 4. Delete only the orphans listed in step 2. A flat file with no sibling +# directory must be kept: deleting it would drop that chromosome entirely. +for f in ./db/variants/*.parquet; do + [ -d "${f%.parquet}" ] && rm "$f" +done + +# 5. Confirm the errors are gone. +afquery check --db ./db/ + +# 6. Re-add the samples with the fixed version. +afquery update-db --db ./db/ --add-samples repair.tsv --bed-dir ./beds/ + +# 7. Verify a variant you know they carry. +afquery variant-info --db ./db/ --locus chr1:887801 +``` + +Re-added samples receive new sample IDs — IDs are never reused after a removal. That is +expected and does not affect results. + +If the original VCFs are no longer available, stop after step 5. The samples are then +absent from both the metadata and `AN`, which is a correct smaller cohort rather than a +biased larger one. Rebuilding with `create-db` is always a valid fallback. + +Databases that were only ever built with `create-db`, and never updated, are unaffected. + +--- + ## Compact Takes a Long Time **Symptom:** `afquery update-db --compact` runs for many minutes or hours. @@ -116,6 +194,7 @@ afquery info --db ./db/ --samples | grep SAMP | `Missing Parquet for chromosome chr3` | Re-run `create-db` or investigate incomplete build | | `Manifest mismatch: expected N samples, found M` | Database may be partially updated; re-run `update-db` | | `Capture file missing for wes_v1` | BED file was not provided at build time; rebuild with `--bed-dir` | +| `chr1: both variants/chr1/ ... and variants/chr1.parquet exist` | Samples added by a pre-0.4.1 `update-db` are invisible to queries; see [Samples Added by update-db Are Missing From Queries](#samples-added-by-update-db-are-missing-from-queries) | --- diff --git a/src/afquery/annotate.py b/src/afquery/annotate.py index 707e256..48e501f 100644 --- a/src/afquery/annotate.py +++ b/src/afquery/annotate.py @@ -3,6 +3,7 @@ import warnings import duckdb +from . import storage from .bitmaps import deserialize from .constants import normalize_chrom from .models import AfqueryWarning, SampleFilter @@ -48,12 +49,12 @@ def _compute_chunk_annotations( n_bitmap_cols = 5 if engine._has_coverage_data else 3 variant_data: dict[tuple[int, str, str], tuple] = {} _db = Path(db_path) - bucket_start = bucket_id * 1_000_000 - bucket_end = (bucket_id + 1) * 1_000_000 - 1 + bucket_start = bucket_id * storage.BUCKET_SIZE + bucket_end = (bucket_id + 1) * storage.BUCKET_SIZE - 1 cols = ", ".join(engine._bitmap_cols(with_pos=True)) - if chrom in engine._partitioned_chroms: - parquet_file = _db / "variants" / chrom / f"bucket_{bucket_id}.parquet" + if storage.chrom_layout(_db / "variants", chrom) == storage.PARTITIONED: + parquet_file = storage.bucket_path(_db / "variants", chrom, bucket_id) if valid_positions and parquet_file.exists(): con = duckdb.connect() placeholders = ", ".join("?" * len(valid_positions)) @@ -67,7 +68,7 @@ def _compute_chunk_annotations( pos, ref, alt = row[0], row[1], row[2] variant_data[(pos, ref, alt)] = tuple(bytes(b) for b in row[3:3 + n_bitmap_cols]) else: - parquet_file = _db / "variants" / f"{chrom}.parquet" + parquet_file = storage.flat_path(_db / "variants", chrom) if valid_positions and parquet_file.exists(): con = duckdb.connect() rows = con.execute( @@ -185,7 +186,7 @@ def annotate_vcf( for variant in vcf: norm = normalize_chrom(variant.CHROM) - bucket = variant.POS // 1_000_000 + bucket = storage.bucket_id(variant.POS) key = (norm, bucket) if key not in variant_buffers: work_order.append(key) diff --git a/src/afquery/benchmark.py b/src/afquery/benchmark.py index 06fb713..72b6bd1 100644 --- a/src/afquery/benchmark.py +++ b/src/afquery/benchmark.py @@ -6,6 +6,7 @@ import pyarrow.parquet as pq +from . import storage from .database import Database @@ -116,32 +117,18 @@ def _find_test_variants( return [] results: list[tuple[str, int, str, str]] = [] - for entry in sorted(variants_dir.iterdir()): + for entry in storage.iter_variant_parquets(variants_dir): if len(results) >= n: break - if entry.suffix == ".parquet": - chrom = entry.stem - tbl = pq.read_table(str(entry), columns=["pos", "ref", "alt"]) - for row in range(min(len(tbl), n - len(results))): - results.append(( - chrom, - int(tbl["pos"][row].as_py()), - str(tbl["ref"][row].as_py()), - str(tbl["alt"][row].as_py()), - )) - elif entry.is_dir(): - chrom = entry.name - for bucket in sorted(entry.glob("bucket_*.parquet")): - if len(results) >= n: - break - tbl = pq.read_table(str(bucket), columns=["pos", "ref", "alt"]) - for row in range(min(len(tbl), n - len(results))): - results.append(( - chrom, - int(tbl["pos"][row].as_py()), - str(tbl["ref"][row].as_py()), - str(tbl["alt"][row].as_py()), - )) + chrom = entry.stem if entry.parent == variants_dir else entry.parent.name + tbl = pq.read_table(str(entry), columns=["pos", "ref", "alt"]) + for row in range(min(len(tbl), n - len(results))): + results.append(( + chrom, + int(tbl["pos"][row].as_py()), + str(tbl["ref"][row].as_py()), + str(tbl["alt"][row].as_py()), + )) return results diff --git a/src/afquery/dump.py b/src/afquery/dump.py index 880ac64..51262df 100644 --- a/src/afquery/dump.py +++ b/src/afquery/dump.py @@ -7,6 +7,7 @@ import duckdb +from . import storage from .bitmaps import deserialize from .constants import normalize_chrom, ALL_CHROMS from .models import SampleFilter @@ -14,7 +15,7 @@ logger = logging.getLogger(__name__) -BUCKET_SIZE = 1_000_000 +BUCKET_SIZE = storage.BUCKET_SIZE def _build_groups(engine, base_sf, by_sex, by_tech, by_phenotype, all_groups): @@ -136,8 +137,8 @@ def _dump_bucket_worker( bucket_end = (bucket_id + 1) * BUCKET_SIZE - 1 # Resolve parquet path and WHERE clause - if chrom in engine._partitioned_chroms: - parquet_file = _db / "variants" / chrom / f"bucket_{bucket_id}.parquet" + if storage.chrom_layout(_db / "variants", chrom) == storage.PARTITIONED: + parquet_file = storage.bucket_path(_db / "variants", chrom, bucket_id) if not parquet_file.exists(): return [] where_parts = [] @@ -150,7 +151,7 @@ def _dump_bucket_worker( params.append(pos_end) where_clause = ("WHERE " + " AND ".join(where_parts)) if where_parts else "" else: - parquet_file = _db / "variants" / f"{chrom}.parquet" + parquet_file = storage.flat_path(_db / "variants", chrom) if not parquet_file.exists(): return [] range_start = max(bucket_start, pos_start) if pos_start is not None else bucket_start @@ -326,23 +327,15 @@ def dump_database( # All chroms that have data available = set() for chrom in ALL_CHROMS: - if chrom in engine._partitioned_chroms: - available.add(chrom) - elif (variants_dir / f"{chrom}.parquet").exists(): + if storage.variant_parquet_glob(variants_dir, chrom) is not None: available.add(chrom) chroms = [c for c in ALL_CHROMS if c in available] # Build work units: (chrom, bucket_id) in genomic order work_units: list[tuple[str, int]] = [] for chrom in chroms: - if chrom in engine._partitioned_chroms: - chrom_dir = variants_dir / chrom - bucket_files = sorted( - chrom_dir.glob("bucket_*.parquet"), - key=lambda p: int(p.stem.split("_")[1]), - ) - for bf in bucket_files: - bid = int(bf.stem.split("_")[1]) + if storage.chrom_layout(variants_dir, chrom) == storage.PARTITIONED: + for bid in storage.existing_bucket_ids(variants_dir, chrom): # Filter by region if specified if pos_start is not None and (bid + 1) * BUCKET_SIZE - 1 < pos_start: continue @@ -350,7 +343,7 @@ def dump_database( continue work_units.append((chrom, bid)) else: - flat_path = variants_dir / f"{chrom}.parquet" + flat_path = storage.flat_path(variants_dir, chrom) if not flat_path.exists(): continue bucket_ids = _discover_flat_buckets(flat_path, pos_start, pos_end) diff --git a/src/afquery/preprocess/build.py b/src/afquery/preprocess/build.py index 507598e..5c03d6a 100644 --- a/src/afquery/preprocess/build.py +++ b/src/afquery/preprocess/build.py @@ -14,12 +14,13 @@ import pyarrow.parquet as pq from pyroaring import BitMap +from .. import storage from ..bitmaps import serialize from ..constants import ALL_CHROMS logger = logging.getLogger(__name__) -BUCKET_SIZE = 1_000_000 +BUCKET_SIZE = storage.BUCKET_SIZE PARQUET_SCHEMA = pa.schema([ ("pos", pa.uint32()), @@ -685,11 +686,9 @@ def build_all_parquets( for chrom in valid_chroms: if resume: if partitioned: - chrom_dir = os.path.join(variants_dir, chrom) - done = (os.path.isdir(chrom_dir) and - bool(glob_module.glob(os.path.join(chrom_dir, "bucket_*.parquet")))) + done = bool(storage.existing_bucket_ids(variants_dir, chrom)) else: - done = os.path.exists(os.path.join(variants_dir, f"{chrom}.parquet")) + done = storage.flat_path(variants_dir, chrom).exists() if done: skipped_chroms.append(chrom) continue diff --git a/src/afquery/preprocess/compact.py b/src/afquery/preprocess/compact.py index 5345b0b..618ec59 100644 --- a/src/afquery/preprocess/compact.py +++ b/src/afquery/preprocess/compact.py @@ -1,4 +1,3 @@ -import glob as glob_module import json import logging import os @@ -11,6 +10,7 @@ import pyarrow.parquet as pq from pyroaring import BitMap +from .. import storage from ..bitmaps import deserialize, serialize from .build import PARQUET_SCHEMA @@ -43,13 +43,7 @@ def compact_database(db_path: Path) -> dict: active_ids = BitMap([r[0] for r in rows]) # Collect all parquet files (flat + partitioned buckets) - all_parquets: list[Path] = [] - for f in sorted(variants_dir.glob("*.parquet")): - all_parquets.append(f) - for chrom_dir in sorted(variants_dir.iterdir()): - if chrom_dir.is_dir(): - for f in sorted(chrom_dir.glob("bucket_*.parquet")): - all_parquets.append(f) + all_parquets: list[Path] = list(storage.iter_variant_parquets(variants_dir)) logger.info("[compact] Compacting %d Parquet file(s) against %d active sample(s)...", len(all_parquets), len(active_ids)) @@ -115,8 +109,11 @@ def compact_database(db_path: Path) -> dict: logger.debug(" [compact] %s: no changes", parquet_file.name) continue - # Build new table with kept rows and updated bitmaps - orig_keep = table.take(keep_indices) + # Build new table with kept rows and updated bitmaps. + # The index type is spelled out: a bare [] makes pyarrow infer a null + # array, which has no take kernel, and every row of a file can legitimately + # be dropped when the removed samples were the only carriers in it. + orig_keep = table.take(pa.array(keep_indices, type=pa.int64())) new_table = pa.table( { "pos": orig_keep["pos"], diff --git a/src/afquery/preprocess/update.py b/src/afquery/preprocess/update.py index 1a8144b..b36e49f 100644 --- a/src/afquery/preprocess/update.py +++ b/src/afquery/preprocess/update.py @@ -7,17 +7,17 @@ import time from dataclasses import dataclass from datetime import datetime, timezone -from pathlib import Path import duckdb import pyarrow as pa import pyarrow.parquet as pq from pyroaring import BitMap +from .. import storage from ..bitmaps import deserialize, serialize, build_phenotype_bitmaps, build_sex_bitmaps, build_tech_bitmaps from ..constants import VALID_GENOME_BUILDS from ..models import Sample, Technology -from .build import PARQUET_SCHEMA, get_chroms_in_temp_files +from .build import PARQUET_SCHEMA, _make_table, _passes_quality, get_chroms_in_temp_files from .ingest import ingest_all from .manifest import parse_manifest from .regions import build_capture_indices @@ -117,26 +117,28 @@ def _regenerate_precomputed_bitmaps(con: sqlite3.Connection) -> None: ) -def _merge_chromosome_parquet( - chrom: str, - db_dir: str, - update_tmp_dir: str, - row_group_size: int = 100_000, - coverage_filter: dict | None = None, - wes_tech_bitmaps: dict[int, BitMap] | None = None, +def _merge_rows_into_file( + out_path: str, + new_rows: list[tuple[int, str, str, BitMap, BitMap, BitMap, BitMap]], + has_phase2: bool, + min_covered: int, + wes_tech_bitmaps: dict[int, BitMap] | None, + row_group_size: int, ) -> tuple[int, int]: - """Merge new temp files into existing chrom Parquet. Returns (new_variants, updated_variants). + """Merge pre-aggregated rows into one Parquet file, rewriting it atomically. - For Phase 2 DBs (coverage_filter active): also merges quality_pass_bitmap and recomputes - filtered_bitmap per row using manifest thresholds and current WES tech bitmaps. - """ - variants_dir = os.path.join(db_dir, "variants") - out_path = os.path.join(variants_dir, f"{chrom}.parquet") + `out_path` is a flat per-chromosome file or a single bucket; the read, merge, + Phase 2 recompute, sort and write are identical either way. + `new_rows` may be empty. That is the recompute-only call used for buckets of + an affected chromosome that received no new data but whose filtered_bitmap + still depends on the enlarged cohort. Such a file is rewritten only when the + recomputation actually moves a bitmap. + + Returns (new_variants, updated_variants). + """ # Read existing Parquet via pyarrow (NOT DuckDB — need Python bitmap deserialization) existing: dict[tuple, tuple[BitMap, BitMap, BitMap, BitMap, BitMap]] = {} - existing_has_fail = False - existing_has_phase2 = False if os.path.exists(out_path): table = pq.read_table(out_path) existing_has_fail = "fail_bitmap" in table.schema.names @@ -159,6 +161,86 @@ def _merge_chromosome_parquet( qp_bm = BitMap() existing[(pos, ref, alt)] = (het_bm, hom_bm, fail_bm, filt_bm, qp_bm) + dirty = bool(new_rows) + new_variants = 0 + updated_variants = 0 + + for pos, ref, alt, new_het, new_hom, new_fail, new_qp in new_rows: + key = (pos, ref, alt) + if key in existing: + old_het, old_hom, old_fail, old_filt, old_qp = existing[key] + existing[key] = ( + old_het | new_het, + old_hom | new_hom, + old_fail | new_fail, + old_filt, + old_qp | new_qp, + ) + updated_variants += 1 + else: + existing[key] = (new_het, new_hom, new_fail, BitMap(), new_qp) + new_variants += 1 + + # Phase 2: recompute filtered_bitmap per row (all rows, since merging may shift any tech) + if has_phase2 and wes_tech_bitmaps and min_covered > 0: + for key, (het, hom, fail, old_filt, qp) in existing.items(): + carrier_set = het | hom | fail + new_filt = BitMap() + for tech_bm in wes_tech_bitmaps.values(): + qp_count = len(qp & tech_bm) + if qp_count < min_covered: + new_filt |= (tech_bm - carrier_set) + if new_filt != old_filt: + dirty = True + existing[key] = (het, hom, fail, new_filt, qp) + + # Nothing moved: leave the file alone. Only reachable when new_rows was empty. + if not dirty or not existing: + return (0, 0) + + # Sort by (pos, alt) and write atomically + sorted_keys = sorted(existing.keys(), key=lambda k: (k[0], k[2])) + + table = _make_table( + [k[0] for k in sorted_keys], + [k[1] for k in sorted_keys], + [k[2] for k in sorted_keys], + [serialize(existing[k][0]) for k in sorted_keys], + [serialize(existing[k][1]) for k in sorted_keys], + [serialize(existing[k][2]) for k in sorted_keys], + [serialize(existing[k][3]) for k in sorted_keys], + [serialize(existing[k][4]) for k in sorted_keys], + ) + + storage.ensure_parent(out_path) + tmp_path = out_path + ".tmp" + pq.write_table(table, tmp_path, row_group_size=row_group_size) + os.replace(tmp_path, out_path) + + return (new_variants, updated_variants) + + +def _merge_chromosome_parquet( + chrom: str, + db_dir: str, + update_tmp_dir: str, + row_group_size: int = 100_000, + coverage_filter: dict | None = None, + wes_tech_bitmaps: dict[int, BitMap] | None = None, + layout: str | None = None, +) -> tuple[int, int]: + """Merge new temp files into a chromosome's variant store. + + Returns (new_variants, updated_variants). + + Writes into whichever layout the chromosome already uses; `layout` decides + for a chromosome new to the database, and defaults to what create-db + produces. For Phase 2 DBs (coverage_filter active) also merges + quality_pass_bitmap and recomputes filtered_bitmap using manifest thresholds + and current WES tech bitmaps. + """ + variants_dir = os.path.join(db_dir, "variants") + # Check if there are any new temp files parquet_files = glob_module.glob(os.path.join(update_tmp_dir, "sample_*.parquet")) if not parquet_files: @@ -202,10 +284,6 @@ def _merge_chromosome_parquet( finally: con.close() - # No new data for this chrom — skip (don't touch the file) - if not rows: - return (0, 0) - coverage_filter = coverage_filter or {} min_dp = coverage_filter.get("min_dp", 0) min_gq = coverage_filter.get("min_gq", 0) @@ -213,8 +291,16 @@ def _merge_chromosome_parquet( min_covered = coverage_filter.get("min_covered", 0) has_phase2 = (min_dp > 0 or min_gq > 0 or min_qual > 0 or min_covered > 0) - new_variants = 0 - updated_variants = 0 + # filtered_bitmap is (tech_bm - carrier_set) row by row, so enlarging a WES + # tech changes it at every row of the database, not only where the new rows + # landed. A chromosome that received nothing still has to be revisited. + recompute_phase2 = bool(has_phase2 and wes_tech_bitmaps and min_covered > 0) + + # Nothing new here and nothing to recompute — leave the files alone. + if not rows and not recompute_phase2: + return (0, 0) + + prepared: list[tuple[int, str, str, BitMap, BitMap, BitMap, BitMap]] = [] for row in rows: if new_has_quality: @@ -225,17 +311,12 @@ def _merge_chromosome_parquet( gqs = [None] * len(sample_ids) quals = [None] * len(sample_ids) - key = (pos, ref, alt) het_ids = [sid for sid, ac, fp in zip(sample_ids, gt_acs, filter_passes) if ac == 1 and fp] hom_ids = [sid for sid, ac, fp in zip(sample_ids, gt_acs, filter_passes) if ac == 2 and fp] fail_ids = [sid for sid, fp in zip(sample_ids, filter_passes) if not fp] - new_het = BitMap(het_ids) - new_hom = BitMap(hom_ids) - new_fail = BitMap(fail_ids) # New quality_pass_ids (only if Phase 2 active) if has_phase2: - from .build import _passes_quality new_qp_ids = [ sid for sid, ac, fp, dp, gq, qual in zip( sample_ids, gt_acs, filter_passes, dps, gqs, quals @@ -246,61 +327,41 @@ def _merge_chromosome_parquet( else: new_qp = BitMap() - if key in existing: - old_het, old_hom, old_fail, old_filt, old_qp = existing[key] - merged_het = old_het | new_het - merged_hom = old_hom | new_hom - merged_fail = old_fail | new_fail - merged_qp = old_qp | new_qp - existing[key] = (merged_het, merged_hom, merged_fail, old_filt, merged_qp) - updated_variants += 1 - else: - existing[key] = (new_het, new_hom, new_fail, BitMap(), new_qp) - new_variants += 1 + prepared.append( + (pos, ref, alt, BitMap(het_ids), BitMap(hom_ids), BitMap(fail_ids), new_qp) + ) - # Phase 2: recompute filtered_bitmap per row (all rows, since merging may shift any tech) - if has_phase2 and wes_tech_bitmaps and min_covered > 0: - for key, (het, hom, fail, _old_filt, qp) in existing.items(): - carrier_set = het | hom | fail - new_filt = BitMap() - for tech_bm in wes_tech_bitmaps.values(): - qp_count = len(qp & tech_bm) - if qp_count < min_covered: - new_filt |= (tech_bm - carrier_set) - existing[key] = (het, hom, fail, new_filt, qp) + if layout is None: + layout = storage.detect_layout(variants_dir) + chrom_fmt = storage.chrom_layout(variants_dir, chrom, layout) - # Sort by (pos, alt) and write atomically - sorted_keys = sorted(existing.keys(), key=lambda k: (k[0], k[2])) + if chrom_fmt == storage.FLAT: + return _merge_rows_into_file( + str(storage.flat_path(variants_dir, chrom)), prepared, + has_phase2, min_covered, wes_tech_bitmaps, row_group_size, + ) - positions = [k[0] for k in sorted_keys] - refs = [k[1] for k in sorted_keys] - alts = [k[2] for k in sorted_keys] - het_bitmaps = [serialize(existing[k][0]) for k in sorted_keys] - hom_bitmaps = [serialize(existing[k][1]) for k in sorted_keys] - fail_bitmaps = [serialize(existing[k][2]) for k in sorted_keys] - filtered_bitmaps = [serialize(existing[k][3]) for k in sorted_keys] - quality_pass_bitmaps = [serialize(existing[k][4]) for k in sorted_keys] + by_bucket: dict[int, list] = {} + for prepared_row in prepared: + by_bucket.setdefault(storage.bucket_id(prepared_row[0]), []).append(prepared_row) - table = pa.table( - { - "pos": pa.array(positions, type=pa.uint32()), - "ref": pa.array(refs, type=pa.large_utf8()), - "alt": pa.array(alts, type=pa.large_utf8()), - "het_bitmap": pa.array(het_bitmaps, type=pa.large_binary()), - "hom_bitmap": pa.array(hom_bitmaps, type=pa.large_binary()), - "fail_bitmap": pa.array(fail_bitmaps, type=pa.large_binary()), - "filtered_bitmap": pa.array(filtered_bitmaps, type=pa.large_binary()), - "quality_pass_bitmap": pa.array(quality_pass_bitmaps, type=pa.large_binary()), - }, - schema=PARQUET_SCHEMA, - ) + target_ids = set(by_bucket) + if recompute_phase2: + # Buckets with no new data must still be revisited, or their + # N_NO_COVERAGE goes stale and the added samples read as hom-ref there. + target_ids |= set(storage.existing_bucket_ids(variants_dir, chrom)) - os.makedirs(variants_dir, exist_ok=True) - tmp_path = out_path + ".tmp" - pq.write_table(table, tmp_path, row_group_size=row_group_size) - os.replace(tmp_path, out_path) + total_new = 0 + total_updated = 0 + for bid in sorted(target_ids): + n, u = _merge_rows_into_file( + str(storage.bucket_path(variants_dir, chrom, bid)), by_bucket.get(bid, []), + has_phase2, min_covered, wes_tech_bitmaps, row_group_size, + ) + total_new += n + total_updated += u - return (new_variants, updated_variants) + return (total_new, total_updated) def _clear_bits_from_parquet(parquet_file: str, removal_ids: BitMap) -> None: @@ -601,12 +662,26 @@ def add_samples( f"genome_build mismatch: DB has '{db_genome_build}', got '{genome_build}'" ) - # 3. Parse new manifest + # 3. Refuse a database that already stores a chromosome in both layouts. + # Checked before ingest so an operator is not told this after hours of work. + mixed = storage.mixed_layout_chroms(os.path.join(db_dir, "variants")) + if mixed: + raise UpdateError( + "Mixed variant layout: chromosome(s) " + + ", ".join(mixed) + + " have both variants// (bucketed) and variants/.parquet. " + "The flat files were written by an older update-db and are invisible to " + "queries; adding more samples would compound the error. Run " + "'afquery check --db ' and follow the recovery procedure in the " + "troubleshooting guide before retrying." + ) + + # 4. Parse new manifest samples_raw, techs_raw = parse_manifest(manifest_path, bed_dir) logger.info("[add-samples] Adding %d new sample(s)...", len(samples_raw)) - # 4. Open connection + # 5. Open connection db_path = os.path.join(db_dir, "metadata.sqlite") con = sqlite3.connect(db_path) @@ -614,7 +689,7 @@ def add_samples( total_updated = 0 try: - # 5. Check for duplicate sample names + # 6. Check for duplicate sample names existing_names = { r[0] for r in con.execute("SELECT sample_name FROM samples").fetchall() } @@ -625,7 +700,7 @@ def add_samples( f"Sample(s) already in database: {', '.join(duplicates)}" ) - # 6. Assign IDs sequentially from next available. + # 7. Assign IDs sequentially from next available. # Use manifest's next_sample_id when present (survives removals) so IDs # are never reused even after a sample has been deleted. manifest_next = manifest.get("next_sample_id") @@ -641,7 +716,7 @@ def add_samples( for i, ps in enumerate(samples_raw) ] - # 7. Handle technologies + # 8. Handle technologies capture_dir = os.path.join(db_dir, "capture") os.makedirs(capture_dir, exist_ok=True) @@ -677,7 +752,7 @@ def add_samples( vcf_paths = [ps.vcf_path for ps in samples_raw] - # 8. Ingest VCFs into fresh tmp_dir + # 9. Ingest VCFs into fresh tmp_dir auto_tmp = tmp_dir is None if auto_tmp: tmp_dir = tempfile.mkdtemp(prefix="afquery_update_") @@ -685,13 +760,14 @@ def add_samples( try: ingest_all(new_samples, vcf_paths, tmp_dir, n_workers=effective_threads) - # 9. Collect chroms from new temp files + # 10. Collect chroms from new temp files chroms = get_chroms_in_temp_files(tmp_dir) # Phase 2: build WES tech bitmaps from CURRENT DB state (existing + new samples) # so that filtered_bitmap recomputation uses the merged cohort. coverage_filter = manifest.get("coverage_filter", {}) wes_tech_bitmaps: dict[int, BitMap] = {} + new_wes_sample = False if coverage_filter and coverage_filter.get("min_covered", 0) > 0: rows = con.execute( "SELECT s.sample_id, s.tech_id, t.bed_path FROM samples s" @@ -708,14 +784,29 @@ def add_samples( None, ) if tech_obj is not None and tech_obj.bed_path is not None: + new_wes_sample = True wes_tech_bitmaps.setdefault(s.tech_id, BitMap()).add(s.sample_id) - # 10. Merge Parquet files + # 11. Merge Parquet files. The layout is decided once for the whole + # update so every chromosome in this batch agrees, including any + # chromosome new to the database. + variants_dir = os.path.join(db_dir, "variants") + if wes_tech_bitmaps and new_wes_sample: + # A WES tech just grew, and filtered_bitmap is derived from the + # tech bitmaps at every row of every chromosome. Restricting the + # recomputation to the chromosomes in this batch would leave the + # added samples counted as homozygous reference everywhere else + # their capture BED reaches. No tech grows for a WGS-only batch, + # so nothing off the touched chromosomes can move and the whole + # store does not need reading. + chroms = sorted(set(chroms) | storage.stored_chroms(variants_dir)) + layout = storage.detect_layout(variants_dir) for chrom in chroms: n, u = _merge_chromosome_parquet( chrom, db_dir, tmp_dir, coverage_filter=coverage_filter, wes_tech_bitmaps=wes_tech_bitmaps or None, + layout=layout, ) total_new += n total_updated += u @@ -725,7 +816,7 @@ def add_samples( import shutil shutil.rmtree(tmp_dir, ignore_errors=True) - # 11. Insert new samples and Phenotype pairs into SQLite + # 12. Insert new samples and Phenotype pairs into SQLite ingested_at = datetime.now(timezone.utc).isoformat() con.executemany( "INSERT INTO samples (sample_id, sample_name, sex, tech_id, vcf_path, ingested_at)" @@ -740,10 +831,10 @@ def add_samples( ] con.executemany("INSERT INTO sample_phenotype VALUES (?, ?)", sample_phenotype_pairs) - # 12. Regenerate precomputed bitmaps + # 13. Regenerate precomputed bitmaps _regenerate_precomputed_bitmaps(con) - # 13. Append changelog entry + # 14. Append changelog entry import json as _json sample_names_json = _json.dumps([s.sample_name for s in new_samples]) con.execute( @@ -752,15 +843,19 @@ def add_samples( ) con.commit() + # sample_count is the number of samples the database holds, which is not + # the next free id once anything has been removed. + sample_count = con.execute("SELECT COUNT(*) FROM samples").fetchone()[0] + finally: con.close() - # 14. Update manifest (persist next_sample_id so future adds don't reuse IDs) + # 15. Update manifest (persist next_sample_id so future adds don't reuse IDs) next_id = starting_id + len(new_samples) # Resolve db_version: explicit value overrides auto-bump current_version = manifest.get("db_version", "1.0") new_version = db_version if db_version is not None else _bump_version(current_version) - _update_manifest(db_dir, next_id, next_sample_id=next_id, db_version=new_version) + _update_manifest(db_dir, sample_count, next_sample_id=next_id, db_version=new_version) logger.info("[add-samples] Done. %d sample(s) added.", len(new_samples)) @@ -799,22 +894,11 @@ def remove_samples(db_dir: str, sample_names: list[str]) -> dict: removal_ids = BitMap([r[0] for r in rows]) id_list = list(removal_ids) - # 2-3. Clear bits from all Parquet files (flat and partitioned) + # 2-3. Clear bits from every Parquet file, in either layout variants_dir = os.path.join(db_dir, "variants") - if os.path.exists(variants_dir): - for pq_file in sorted( - glob_module.glob(os.path.join(variants_dir, "*.parquet")) - ): - _clear_bits_from_parquet(pq_file, removal_ids) - logger.debug(" [remove-samples] %s cleared", os.path.basename(pq_file)) - # Also handle partitioned format (variants/{chrom}/bucket_*.parquet) - for chrom_dir in sorted(Path(variants_dir).iterdir()): - if chrom_dir.is_dir(): - for pq_file in sorted( - glob_module.glob(str(chrom_dir / "bucket_*.parquet")) - ): - _clear_bits_from_parquet(pq_file, removal_ids) - logger.debug(" [remove-samples] %s cleared", os.path.basename(pq_file)) + for pq_file in storage.iter_variant_parquets(variants_dir): + _clear_bits_from_parquet(str(pq_file), removal_ids) + logger.debug(" [remove-samples] %s cleared", pq_file.name) # 4-5. Delete from SQLite ph2 = ",".join("?" * len(id_list)) @@ -950,20 +1034,20 @@ def info(msg: str) -> None: err("variants/ directory not found") return results - # Collect flat parquet files - flat_parquets = sorted(glob_module.glob(os.path.join(variants_dir, "*.parquet"))) - # Collect partitioned parquets (variants/{chrom}/bucket_*.parquet) - chrom_dirs = [] - bucket_parquets = [] - for entry in sorted(os.scandir(variants_dir), key=lambda e: e.name): - if entry.is_dir(): - chrom_dirs.append(entry.path) - bucket_parquets.extend( - sorted(glob_module.glob(os.path.join(entry.path, "bucket_*.parquet"))) - ) + # Check 9b: no chromosome is stored in both layouts at once + for chrom in storage.mixed_layout_chroms(variants_dir): + err( + f"{chrom}: both variants/{chrom}/ (bucketed) and variants/{chrom}.parquet " + f"exist. Queries read only the bucketed files, so every sample whose calls " + f"live in {chrom}.parquet is silently counted as homozygous reference and " + f"biases allele frequencies downward. See the troubleshooting guide: " + f"'Samples Added by update-db Are Missing From Queries'." + ) - parquet_files = flat_parquets + bucket_parquets - n_chroms = len(flat_parquets) + len(chrom_dirs) + parquet_files = [str(p) for p in storage.iter_variant_parquets(variants_dir)] + n_chroms = len(storage.flat_chroms(variants_dir)) + len( + storage.partitioned_chroms(variants_dir) + ) expected_fields = { "pos": pa.uint32(), diff --git a/src/afquery/query.py b/src/afquery/query.py index 8d17665..4939c30 100644 --- a/src/afquery/query.py +++ b/src/afquery/query.py @@ -6,6 +6,7 @@ import duckdb from pyroaring import BitMap +from . import storage from .bitmaps import deserialize from .capture import CaptureIndex, describe_capture_problem, load_capture_indices from .constants import normalize_chrom, ALL_CHROMS, CHROM_ORDER @@ -83,18 +84,9 @@ def __init__(self, db_path: str): } # Cache which chroms use partitioned storage (variants/{chrom}/ directory) - self._partitioned_chroms: set[str] = set() - variants_dir = self._db / "variants" - if variants_dir.exists(): - for p in variants_dir.iterdir(): - if p.is_dir(): - self._partitioned_chroms.add(p.name) - - self._flat_chroms: set[str] = set() - if variants_dir.exists(): - for p in variants_dir.iterdir(): - if p.is_file() and p.suffix == ".parquet": - self._flat_chroms.add(p.stem) + self._variants_dir = self._db / "variants" + self._partitioned_chroms: set[str] = storage.partitioned_chroms(self._variants_dir) + self._flat_chroms: set[str] = storage.flat_chroms(self._variants_dir) self._all_known_chroms: set[str] = self._partitioned_chroms | self._flat_chroms # Precompute covered bitmap for positions where ALL technologies cover (common case) @@ -302,6 +294,19 @@ def _compute_eligible( if capture_idx.covers(chrom, pos): covered |= self._tech_bitmaps.get(str(tech_id), BitMap()) eligible = sample_bitmap & covered + + # A sample with no allele at this position is not eligible at it. The + # case is chrY in a female: she is neither a carrier nor homozygous + # reference there, because she has no chrY to genotype. AN has always + # excluded such samples, so leaving them in the eligible set made + # N_HOM_REF and n_samples_eligible disagree with it. Every other + # chromosome is unaffected: on chrX both sexes carry alleles, only the + # ploidy differs. + haploid, diploid = split_ploidy( + eligible, self._male_bm, self._female_bm, chrom, pos, self._genome_build, + ) + eligible = haploid | diploid + AN = compute_AN( eligible, self._male_bm, @@ -312,23 +317,13 @@ def _compute_eligible( def _parquet_path(self, chrom: str, pos: int) -> Path | None: """Resolve Parquet path for a point query. Partitioned takes priority over flat.""" - if chrom in self._partitioned_chroms: - bucket = pos // 1_000_000 - p = self._db / "variants" / chrom / f"bucket_{bucket}.parquet" - return p if p.exists() else None - flat = self._db / "variants" / f"{chrom}.parquet" - return flat if flat.exists() else None + return storage.variant_parquet_for_pos(self._variants_dir, chrom, pos) def _parquet_glob(self, chrom: str) -> str | None: """Return path/glob pattern for batch/region queries. None if no data for chrom.""" if chrom in self._glob_cache: return self._glob_cache[chrom] - if chrom in self._partitioned_chroms: - chrom_dir = self._db / "variants" / chrom - result = str(chrom_dir / "bucket_*.parquet") if any(chrom_dir.glob("bucket_*.parquet")) else None - else: - flat = self._db / "variants" / f"{chrom}.parquet" - result = str(flat) if flat.exists() else None + result = storage.variant_parquet_glob(self._variants_dir, chrom) self._glob_cache[chrom] = result return result diff --git a/src/afquery/storage.py b/src/afquery/storage.py new file mode 100644 index 0000000..9a2b708 --- /dev/null +++ b/src/afquery/storage.py @@ -0,0 +1,177 @@ +"""Resolution of on-disk paths in the variant Parquet store. + +Two layouts exist. The **partitioned** layout stores one file per 1 Mbp bucket +under a per-chromosome directory:: + + variants/chr1/bucket_0.parquet # positions 0 .. 999,999 + variants/chr1/bucket_1.parquet # positions 1,000,000 .. 1,999,999 + +The **flat** layout stores one file per chromosome:: + + variants/chr1.parquet + +``create-db`` has produced the partitioned layout since it became the default, +so every database built by the normal path is bucketed. The flat layout survives +in small hand-built databases and in test fixtures, and is still readable. + +Readers have always preferred the bucket directory when a chromosome somehow has +both. That rule used to be reimplemented in every module that opened a Parquet +file — the query engine, the dump and annotate workers, the benchmark helper, +the compactor and the updater — which is how the updater came to disagree with +the readers and write rows into a file no query would ever open. Keeping the +rule in one module is what stops writers and readers from drifting apart again. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Iterator + +BUCKET_SIZE = 1_000_000 + +PARTITIONED = "partitioned" +FLAT = "flat" + + +def bucket_id(pos: int) -> int: + """Bucket owning `pos`.""" + return pos // BUCKET_SIZE + + +def bucket_path(variants_dir: Path | str, chrom: str, bucket: int) -> Path: + """Path of one bucket file, whether or not it exists.""" + return Path(variants_dir) / chrom / f"bucket_{bucket}.parquet" + + +def flat_path(variants_dir: Path | str, chrom: str) -> Path: + """Path of a chromosome's flat file, whether or not it exists.""" + return Path(variants_dir) / f"{chrom}.parquet" + + +def partitioned_chroms(variants_dir: Path | str) -> set[str]: + """Chromosomes stored as a bucket directory. + + Every subdirectory counts, with no check against the canonical chromosome + list: a database may hold unplaced or alt contigs whose names normalize to + bodies like 'chrGL000209.1', and filtering them here would make their data + unreadable. + """ + variants_dir = Path(variants_dir) + if not variants_dir.exists(): + return set() + return {p.name for p in variants_dir.iterdir() if p.is_dir()} + + +def flat_chroms(variants_dir: Path | str) -> set[str]: + """Chromosomes stored as a single flat Parquet file.""" + variants_dir = Path(variants_dir) + if not variants_dir.exists(): + return set() + return {p.stem for p in variants_dir.iterdir() if p.is_file() and p.suffix == ".parquet"} + + +def stored_chroms(variants_dir: Path | str) -> set[str]: + """Every chromosome the database holds, in either layout.""" + return partitioned_chroms(variants_dir) | flat_chroms(variants_dir) + + +def detect_layout(variants_dir: Path | str) -> str: + """Layout of the database as a whole. + + A subdirectory only votes 'partitioned' if it actually holds a bucket file, + so a stray empty directory cannot flip the answer. An empty or missing + variants directory answers 'partitioned', which is what create-db produces — + that way a database with no variants yet grows buckets instead of sprouting + a second layout the readers would ignore. + """ + variants_dir = Path(variants_dir) + if variants_dir.is_dir(): + for entry in variants_dir.iterdir(): + if entry.is_dir() and any(entry.glob("bucket_*.parquet")): + return PARTITIONED + if any(p.suffix == ".parquet" for p in variants_dir.iterdir() if p.is_file()): + return FLAT + return PARTITIONED + + +def chrom_layout(variants_dir: Path | str, chrom: str, default: str = PARTITIONED) -> str: + """Layout of one chromosome; `default` decides for a chromosome new to the database. + + The bucket directory wins when both are present, matching what the readers do. + """ + variants_dir = Path(variants_dir) + if (variants_dir / chrom).is_dir(): + return PARTITIONED + if flat_path(variants_dir, chrom).exists(): + return FLAT + return default + + +def mixed_layout_chroms(variants_dir: Path | str) -> list[str]: + """Chromosomes that have both a bucket directory and a flat file. + + A database should never be in this state. When it is, the flat file is dead + weight: queries read only the bucket directory, so any sample whose calls + live in the flat file is counted as homozygous reference everywhere. + """ + variants_dir = Path(variants_dir) + if not variants_dir.is_dir(): + return [] + return sorted( + p.name for p in variants_dir.iterdir() + if p.is_dir() and flat_path(variants_dir, p.name).exists() + ) + + +def existing_bucket_ids(variants_dir: Path | str, chrom: str) -> list[int]: + """Bucket ids already written for `chrom`, ascending. Empty when not bucketed.""" + # The chromosome is a directory name, not part of the pattern: contig names + # can legitimately contain glob metacharacters (GRCh38 spells HLA contigs + # HLA-A*01:01:01:01), and this is reached with names read straight off disk. + chrom_dir = Path(variants_dir) / chrom + if not chrom_dir.is_dir(): + return [] + ids: list[int] = [] + for p in chrom_dir.glob("bucket_*.parquet"): + stem = p.stem[len("bucket_"):] + if stem.isdigit(): + ids.append(int(stem)) + return sorted(ids) + + +def variant_parquet_for_pos(variants_dir: Path | str, chrom: str, pos: int) -> Path | None: + """File holding `pos` for a point query, or None when there is no data.""" + variants_dir = Path(variants_dir) + if (variants_dir / chrom).is_dir(): + p = bucket_path(variants_dir, chrom, bucket_id(pos)) + return p if p.exists() else None + flat = flat_path(variants_dir, chrom) + return flat if flat.exists() else None + + +def variant_parquet_glob(variants_dir: Path | str, chrom: str) -> str | None: + """Path or glob pattern covering a whole chromosome, or None when there is no data.""" + variants_dir = Path(variants_dir) + chrom_dir = variants_dir / chrom + if chrom_dir.is_dir(): + return str(chrom_dir / "bucket_*.parquet") if any( + chrom_dir.glob("bucket_*.parquet") + ) else None + flat = flat_path(variants_dir, chrom) + return str(flat) if flat.exists() else None + + +def iter_variant_parquets(variants_dir: Path | str) -> Iterator[Path]: + """Every variant Parquet file in the store, flat files first, then buckets.""" + variants_dir = Path(variants_dir) + if not variants_dir.is_dir(): + return + yield from sorted(variants_dir.glob("*.parquet")) + for chrom_dir in sorted(p for p in variants_dir.iterdir() if p.is_dir()): + yield from sorted(chrom_dir.glob("bucket_*.parquet")) + + +def ensure_parent(path: Path | str) -> None: + """Create the directory holding `path`, including a per-chromosome bucket directory.""" + os.makedirs(os.path.dirname(str(path)), exist_ok=True) diff --git a/tests/data/expected_results.json b/tests/data/expected_results.json index 2b09a73..c9543cd 100644 --- a/tests/data/expected_results.json +++ b/tests/data/expected_results.json @@ -45,7 +45,7 @@ "n_eligible": 3 }, { - "description": "chrY:500000 J45 both sexes (only males count for AN)", + "description": "chrY:500000 J45 both sexes (females have no chrY to genotype)", "chrom": "chrY", "pos": 500000, "ref": "T", @@ -57,7 +57,7 @@ "AC": 1, "AN": 1, "AF": 1.0, - "n_eligible": 2 + "n_eligible": 1 }, { "description": "chrM:100 E11.9 both sexes (ploidy=1 for all)", diff --git a/tests/oracle.py b/tests/oracle.py new file mode 100644 index 0000000..5cfc614 --- /dev/null +++ b/tests/oracle.py @@ -0,0 +1,195 @@ +"""An independent reimplementation of what a query is supposed to return. + +This module deliberately imports nothing from afquery. It reads the same raw +inputs the database was built from — VCF text, BED text, the manifest TSV — and +works out the expected counts from first principles. Reusing the code under test +to predict its own output would only assert that the code is self-consistent, +which is exactly what the existing suite already did while add-samples was +silently dropping samples on the floor. + +The coverage half is the half that matters. AC, N_HET and N_HOM_ALT come from +the VCFs, but AN, N_HOM_REF and N_NO_COVERAGE come from capture-BED coverage +crossed with sex-dependent ploidy, and every silent counting bug this project +has had lived on that side. +""" + +from __future__ import annotations + +from pathlib import Path + +# PAR coordinates, 1-based inclusive, written out here rather than imported. +PAR = { + "GRCh37": { + "chrX": [(60_001, 2_699_520), (154_931_044, 155_260_560)], + "chrY": [(10_001, 2_649_520), (59_034_050, 59_363_566)], + }, + "GRCh38": { + "chrX": [(10_001, 2_781_479), (155_701_383, 156_030_895)], + "chrY": [(10_001, 2_781_479), (56_887_903, 57_217_415)], + }, +} + + +def _canon(chrom: str) -> str: + """Fold the naming conventions a BED or VCF may use onto one spelling.""" + c = chrom.strip() + if c[:3].lower() == "chr": + c = c[3:] + if c.upper() in ("M", "MT"): + return "chrM" + if c.upper() in ("X", "Y"): + return "chr" + c.upper() + return "chr" + c + + +class Cohort: + """Samples, their genotypes and their capture regions, read from raw files.""" + + def __init__(self, manifest_path, bed_dir=None, genome_build="GRCh37"): + self.genome_build = genome_build + self.samples: list[str] = [] + self.sex: dict[str, str] = {} + self.tech: dict[str, str] = {} + self.phenotypes: dict[str, list[str]] = {} + # (chrom, pos, ref, alt) -> {sample: "het" | "hom" | "fail"} + self.calls: dict[tuple[str, int, str, str], dict[str, str]] = {} + # tech -> list of (chrom, start_1based, end_1based); None means whole genome + self.regions: dict[str, list[tuple[str, int, int]] | None] = {} + + manifest_path = Path(manifest_path) + base = manifest_path.parent + rows = _read_tsv(manifest_path) + for row in rows: + name = row["sample_name"] + self.samples.append(name) + self.sex[name] = row["sex"] + self.tech[name] = row["tech_name"] + self.phenotypes[name] = [ + c.strip() for c in row.get("phenotype_codes", "").split(",") if c.strip() + ] + vcf = Path(row["vcf_path"]) + if not vcf.is_absolute(): + vcf = base / vcf + self._read_vcf(vcf, name) + + for tech in set(self.tech.values()): + bed = None if bed_dir is None else Path(bed_dir) / f"{tech}.bed" + if bed is not None and bed.exists(): + self.regions[tech] = _read_bed(bed) + else: + self.regions[tech] = None # no BED: covered everywhere + + def _read_vcf(self, path: Path, sample: str) -> None: + for line in Path(path).read_text().splitlines(): + if not line or line.startswith("#"): + continue + f = line.split("\t") + chrom, pos, ref, alt, filt, fmt, call = ( + _canon(f[0]), int(f[1]), f[3], f[4], f[6], f[8], f[9] + ) + gt = call.split(":")[fmt.split(":").index("GT")] + alleles = [a for a in gt.replace("|", "/").split("/") if a != "."] + passed = filt in ("PASS", ".", "") + for i, one_alt in enumerate(alt.split(","), start=1): + if one_alt == "*": + continue + n = alleles.count(str(i)) + key = (chrom, pos, ref, one_alt) + if not passed and not alleles: + self.calls.setdefault(key, {})[sample] = "fail" + elif n == 0: + continue + elif not passed: + self.calls.setdefault(key, {})[sample] = "fail" + else: + self.calls.setdefault(key, {})[sample] = "hom" if n >= 2 else "het" + + def covers(self, sample: str, chrom: str, pos: int) -> bool: + regions = self.regions[self.tech[sample]] + if regions is None: + return True + return any(c == chrom and start <= pos <= end for c, start, end in regions) + + def is_haploid(self, sample: str, chrom: str, pos: int) -> bool: + if chrom == "chrM": + return True + if chrom == "chrY": + return True + if chrom == "chrX" and not self._in_par(chrom, pos): + return self.sex[sample] == "male" + return False + + def _in_par(self, chrom: str, pos: int) -> bool: + return any(s <= pos <= e for s, e in PAR[self.genome_build].get(chrom, [])) + + def eligible(self, chrom: str, pos: int, samples=None) -> list[str]: + pool = self.samples if samples is None else samples + out = [] + for s in pool: + if not self.covers(s, chrom, pos): + continue + # chrY carries no alleles for females, so they are not eligible there + if chrom == "chrY" and self.sex[s] != "male": + continue + out.append(s) + return out + + def expect(self, chrom, pos, ref, alt, samples=None) -> dict: + """Expected query counts for one variant.""" + elig = self.eligible(chrom, pos, samples) + calls = self.calls.get((chrom, pos, ref, alt), {}) + + AN = 0 + AC = 0 + n_het = n_hom_alt = n_fail = 0 + for s in elig: + haploid = self.is_haploid(s, chrom, pos) + AN += 1 if haploid else 2 + state = calls.get(s) + if state == "fail": + n_fail += 1 + elif state == "het": + if haploid: + AC += 1 + n_hom_alt += 1 + else: + AC += 1 + n_het += 1 + elif state == "hom": + if haploid: + AC += 1 + else: + AC += 2 + n_hom_alt += 1 + + return { + "AC": AC, + "AN": AN, + "N_HET": n_het, + "N_HOM_ALT": n_hom_alt, + "N_FAIL": n_fail, + "N_HOM_REF": len(elig) - n_het - n_hom_alt - n_fail, + } + + def variants(self) -> list[tuple[str, int, str, str]]: + return sorted(self.calls) + + def with_phenotype(self, code: str) -> list[str]: + return [s for s in self.samples if code in self.phenotypes[s]] + + +def _read_tsv(path: Path) -> list[dict]: + lines = [ln for ln in Path(path).read_text().splitlines() if ln.strip()] + header = lines[0].split("\t") + return [dict(zip(header, ln.split("\t"))) for ln in lines[1:]] + + +def _read_bed(path: Path) -> list[tuple[str, int, int]]: + """BED is 0-based half-open; convert to 1-based inclusive.""" + out = [] + for line in Path(path).read_text().splitlines(): + if not line.strip() or line.startswith(("#", "track", "browser")): + continue + f = line.split("\t") + out.append((_canon(f[0]), int(f[1]) + 1, int(f[2]))) + return out diff --git a/tests/test_haploid_stats.py b/tests/test_haploid_stats.py index d635097..0cf935a 100644 --- a/tests/test_haploid_stats.py +++ b/tests/test_haploid_stats.py @@ -145,3 +145,33 @@ def test_region_chrM_haploid(self, test_db): region = db.query_region(chrom="chrM", start=100, end=100) assert len(region) == 1 assert region[0].N_HET == 0 + + +class TestChrYEligibility: + """Females hold no chrY, so they are neither carriers nor homozygous reference there.""" + + def test_chrY_excludes_females_from_eligible(self, test_db): + db = Database(test_db) + r = db.query(chrom="chrY", pos=500000)[0] + + # The fixture is 5 male and 5 female WGS/WES samples; only the males + # that the capture covers can be genotyped on chrY. + assert r.n_samples_eligible == r.AN, ( + "n_samples_eligible counts samples with no chrY: it must agree with AN, " + "which has always been male-only on this chromosome" + ) + + def test_chrY_counts_partition_the_eligible_set(self, test_db): + """The per-genotype counts must add up to the eligible set on every chromosome.""" + db = Database(test_db) + for chrom, pos in (("chr1", 1500), ("chrX", 5000000), ("chrY", 500000), ("chrM", 100)): + for r in db.query(chrom=chrom, pos=pos): + total = r.N_HET + r.N_HOM_ALT + r.N_HOM_REF + r.N_FAIL + r.N_NO_COVERAGE + assert total == r.n_samples_eligible, f"{chrom}:{pos} counts do not partition" + + def test_chrY_no_phantom_hom_ref(self, test_db): + """S00 and S01 carry chrY:500000; nobody else eligible is homozygous reference.""" + db = Database(test_db) + r = db.query(chrom="chrY", pos=500000)[0] + assert r.N_HOM_ALT == 2 + assert r.N_HOM_REF == 0 diff --git a/tests/test_invariants.py b/tests/test_invariants.py new file mode 100644 index 0000000..a615ce0 --- /dev/null +++ b/tests/test_invariants.py @@ -0,0 +1,186 @@ +"""Properties that must hold however a database was assembled. + +These compare a database against another database rather than against a stored +expectation, so they keep holding as the fixtures change. The build-equivalence +property is the one that would have caught add-samples writing into a layout +the query engine does not read. +""" +import shutil + +import pytest + +from afquery import storage +from afquery.database import Database +from afquery.preprocess import run_preprocess +from afquery.preprocess.build import build_all_parquets +from afquery.preprocess.ingest import ingest_all +from afquery.models import Sample + +from test_update import write_manifest, write_vcf + +# (name, sex, tech, [(chrom, pos, ref, alt, gt), ...]) +GROUP_A = [ + ("A0", "male", "wgs", [("chr1", 1000, "A", "T", "0/1"), + ("chr1", 2_400_000, "G", "C", "1/1"), + ("chrX", 4_000_000, "T", "A", "0/1")]), + ("A1", "female", "wgs", [("chr1", 1000, "A", "T", "1/1"), + ("chrM", 300, "C", "T", "0/1")]), +] +GROUP_B = [ + ("B0", "female", "wgs", [("chr1", 1000, "A", "T", "0/1"), + ("chr1", 5_600_000, "A", "G", "0/1"), + ("chr2", 700, "T", "C", "1/1")]), + ("B1", "male", "wgs", [("chr1", 2_400_000, "G", "C", "0/1"), + ("chrY", 700_000, "A", "T", "1/1"), + ("chrX", 4_000_000, "T", "A", "1/1")]), +] + +ALL_SITES = sorted({ + (chrom, pos, ref, alt) + for _n, _s, _t, variants in GROUP_A + GROUP_B + for chrom, pos, ref, alt, _gt in variants +}) + + +def _write_group(tmp_path, group): + entries = [] + for name, sex, tech, variants in group: + vcf = tmp_path / f"{name}.vcf" + write_vcf(str(vcf), name, variants) + entries.append((name, sex, tech, str(vcf), "COHORT")) + return entries + + +def _manifest(tmp_path, entries, filename): + path = tmp_path / filename + write_manifest(str(path), entries) + return str(path) + + +def _snapshot(db_dir): + """Every site's counts, keyed by variant — comparable across sample id orders.""" + db = Database(str(db_dir)) + out = {} + for chrom, pos, ref, alt in ALL_SITES: + for r in db.query(chrom=chrom, pos=pos, sex="both"): + out[(chrom, pos, r.variant.ref, r.variant.alt)] = ( + r.AC, r.AN, r.N_HET, r.N_HOM_ALT, r.N_HOM_REF, r.N_FAIL, + ) + return out + + +@pytest.fixture +def groups(tmp_path): + a = _write_group(tmp_path, GROUP_A) + b = _write_group(tmp_path, GROUP_B) + return a, b + + +def test_add_samples_equals_building_the_union(tmp_path, groups): + """build(A + B) and build(A) then add(B) must answer identically. + + Compared at query level rather than bitmap level: the two routes assign + sample ids in a different order, so the bytes legitimately differ while + every answer must not. + """ + a, b = groups + + together = tmp_path / "together" + run_preprocess(manifest_path=_manifest(tmp_path, a + b, "ab.tsv"), + output_dir=str(together), genome_build="GRCh37", threads=1) + + incremental = tmp_path / "incremental" + run_preprocess(manifest_path=_manifest(tmp_path, a, "a.tsv"), + output_dir=str(incremental), genome_build="GRCh37", threads=1) + from afquery.preprocess.update import add_samples + add_samples(str(incremental), _manifest(tmp_path, b, "b.tsv"), threads=1) + + assert _snapshot(incremental) == _snapshot(together) + + +def test_flat_and_partitioned_answer_identically(tmp_path, groups): + """The storage layout must not be observable through any query.""" + a, b = groups + manifest = _manifest(tmp_path, a + b, "all.tsv") + + part = tmp_path / "part" + run_preprocess(manifest_path=manifest, output_dir=str(part), + genome_build="GRCh37", threads=1) + + flat = tmp_path / "flat" + shutil.copytree(part, flat) + shutil.rmtree(flat / "variants") + (flat / "variants").mkdir() + samples = [Sample(i, name, sex, 0) + for i, (name, sex, _t, _v, _p) in enumerate(a + b)] + vcfs = [str(tmp_path / f"{name}.vcf") for name, _s, _t, _v, _p in a + b] + ingest_tmp = tmp_path / "ingest" + ingest_tmp.mkdir() + ingest_all(samples, vcfs, str(ingest_tmp), n_workers=1) + build_all_parquets(str(ingest_tmp), str(flat / "variants"), + n_workers=1, partitioned=False) + + assert storage.detect_layout(part / "variants") == storage.PARTITIONED + assert storage.detect_layout(flat / "variants") == storage.FLAT + assert _snapshot(flat) == _snapshot(part) + + +def test_remove_undoes_add(tmp_path, groups): + """build(A) then add(B) then remove(B) must answer like build(A). + + Removal clears bits but leaves the rows behind — that is what compact is + for — so sites only group B carried survive as AC=0 rows. Every site the + smaller database knows must answer identically, and the leftovers must + carry nobody. + """ + a, b = groups + + alone = tmp_path / "alone" + run_preprocess(manifest_path=_manifest(tmp_path, a, "a_only.tsv"), + output_dir=str(alone), genome_build="GRCh37", threads=1) + before = _snapshot(alone) + + grown = tmp_path / "grown" + run_preprocess(manifest_path=_manifest(tmp_path, a, "a_dup.tsv"), + output_dir=str(grown), genome_build="GRCh37", threads=1) + from afquery.preprocess.update import add_samples, remove_samples + add_samples(str(grown), _manifest(tmp_path, b, "b_dup.tsv"), threads=1) + remove_samples(str(grown), [name for name, _s, _t, _v, _p in b]) + + after = _snapshot(grown) + for key, counts in before.items(): + assert after[key] == counts, f"{key} changed across add-then-remove" + for key, counts in after.items(): + if key not in before: + AC, _AN, n_het, n_hom_alt, _n_hom_ref, n_fail = counts + assert (AC, n_het, n_hom_alt, n_fail) == (0, 0, 0, 0), ( + f"{key} still has carriers after their only samples were removed" + ) + + +def test_compact_preserves_carried_variants_and_is_idempotent(tmp_path, groups): + """Compact may drop rows nobody carries, but must not change any that remain. + + B1 is the only carrier on chrY:700000, so removing it empties that bucket + entirely — the case where compact used to crash on an empty row selection. + """ + a, b = groups + db = tmp_path / "compactme" + run_preprocess(manifest_path=_manifest(tmp_path, a + b, "cab.tsv"), + output_dir=str(db), genome_build="GRCh37", threads=1) + from afquery.preprocess.compact import compact_database + from afquery.preprocess.update import remove_samples + + remove_samples(str(db), ["B1"]) + before = _snapshot(db) + + compact_database(db) + once = _snapshot(db) + + for key, counts in once.items(): + assert counts == before[key], f"{key} changed across compact" + for key in set(before) - set(once): + assert before[key][0] == 0, f"compact dropped {key}, which still had carriers" + + compact_database(db) + assert _snapshot(db) == once diff --git a/tests/test_oracle_consistency.py b/tests/test_oracle_consistency.py new file mode 100644 index 0000000..121237a --- /dev/null +++ b/tests/test_oracle_consistency.py @@ -0,0 +1,164 @@ +"""Query results checked against an independent reimplementation. + +The oracle (tests/oracle.py) derives the expected counts from the VCF, BED and +manifest text without importing afquery. Anything the database and the oracle +disagree about is a real counting error, not a fixture that drifted. +""" +import shutil + +import pytest + +from afquery.database import Database +from afquery.preprocess import run_preprocess +from afquery.preprocess.build import build_all_parquets +from afquery.preprocess.update import add_samples + +import oracle +from test_update import write_manifest, write_vcf + + +def _assert_matches(db_dir, cohort, samples=None, phenotype=None, label=""): + db = Database(str(db_dir)) + checked = 0 + for chrom, pos, ref, alt in cohort.variants(): + want = cohort.expect(chrom, pos, ref, alt, samples) + kwargs = {"chrom": chrom, "pos": pos, "sex": "both"} + if phenotype is not None: + kwargs["phenotype"] = phenotype + got = [r for r in db.query(**kwargs) + if (r.variant.ref, r.variant.alt) == (ref, alt)] + + if want["AN"] == 0: + assert not got or got[0].AC == 0 + continue + + assert got, f"{label} {chrom}:{pos} {ref}>{alt} returned nothing, expected {want}" + r = got[0] + actual = { + "AC": r.AC, "AN": r.AN, "N_HET": r.N_HET, "N_HOM_ALT": r.N_HOM_ALT, + "N_FAIL": r.N_FAIL, "N_HOM_REF": r.N_HOM_REF, + } + assert actual == want, f"{label} {chrom}:{pos} {ref}>{alt}" + checked += 1 + assert checked > 0, f"{label}: oracle checked nothing" + + +# --------------------------------------------------------------------------- +# Fixture cohort, both layouts +# --------------------------------------------------------------------------- + +@pytest.fixture(scope="module") +def fixture_cohort(data_dir): + return oracle.Cohort(data_dir / "manifest.tsv", bed_dir=data_dir / "beds") + + +@pytest.fixture(scope="module") +def built_db(tmp_path_factory, data_dir): + db = tmp_path_factory.mktemp("oracle_db") + run_preprocess( + manifest_path=str(data_dir / "manifest.tsv"), output_dir=str(db), + genome_build="GRCh37", bed_dir=str(data_dir / "beds"), threads=2, + ) + return db + + +def test_partitioned_build_matches_oracle(built_db, fixture_cohort): + _assert_matches(built_db, fixture_cohort, label="partitioned") + + +def test_flat_build_matches_oracle(tmp_path, built_db, fixture_cohort, data_dir): + """The legacy layout must agree with the oracle too, and so with the buckets.""" + import sqlite3 + + from afquery.models import Sample + from afquery.preprocess.ingest import ingest_all + + flat = tmp_path / "flat" + shutil.copytree(built_db, flat) + shutil.rmtree(flat / "variants") + (flat / "variants").mkdir() + + con = sqlite3.connect(flat / "metadata.sqlite") + rows = con.execute("SELECT sample_id, sample_name, sex, tech_id FROM samples").fetchall() + con.close() + samples = [Sample(*r) for r in rows] + vcfs = [str(data_dir / "vcfs" / f"{r[1]}.vcf") for r in rows] + + tmp = tmp_path / "ingest" + tmp.mkdir() + ingest_all(samples, vcfs, str(tmp), n_workers=1) + build_all_parquets(str(tmp), str(flat / "variants"), n_workers=1, partitioned=False) + + _assert_matches(flat, fixture_cohort, label="flat") + + +def test_phenotype_subset_matches_oracle(built_db, fixture_cohort): + subset = fixture_cohort.with_phenotype("E11.9") + assert subset + _assert_matches(built_db, fixture_cohort, samples=subset, + phenotype=["E11.9"], label="phenotype E11.9") + + +# --------------------------------------------------------------------------- +# After add-samples — the case the bug lived in +# --------------------------------------------------------------------------- + +@pytest.fixture +def grown_db(built_db, tmp_path, data_dir): + """Fixture cohort plus one wgs and one wes_kit_a sample, added via update.""" + db = tmp_path / "grown" + shutil.copytree(built_db, db) + + entries = [] + for name, tech, sex, variants in [ + ("S10", "wgs", "male", [("chr1", 1500, "A", "T", "1/1"), + ("chr1", 7000, "A", "T", "0/1"), + ("chrX", 5000000, "A", "G", "0/1")]), + ("S11", "wes_kit_a", "female", [("chr1", 1500, "A", "T", "0/1"), + ("chr1", 1800, "C", "G", "1/1")]), + ]: + vcf = tmp_path / f"{name}.vcf" + write_vcf(str(vcf), name, variants) + entries.append((name, sex, tech, str(vcf), "ZZNEW")) + + manifest = tmp_path / "grow.tsv" + write_manifest(str(manifest), entries) + add_samples(str(db), str(manifest), threads=1, bed_dir=str(data_dir / "beds")) + return db, manifest + + +def test_after_add_samples_matches_oracle(grown_db, data_dir, tmp_path): + db, grow_manifest = grown_db + + # Oracle over the union of the original manifest and the added one + combined = tmp_path / "combined.tsv" + original = (data_dir / "manifest.tsv").read_text().rstrip("\n").splitlines() + added = grow_manifest.read_text().rstrip("\n").splitlines()[1:] + header = original[0].split("\t") + + out = [original[0]] + for line in original[1:]: + row = dict(zip(header, line.split("\t"))) + # the fixture manifest's paths are relative to tests/data + row["vcf_path"] = str(data_dir / row["vcf_path"]) + out.append("\t".join(row[c] for c in header)) + for line in added: + # write_manifest emits a different column order + row = dict(zip(["sample_name", "sex", "tech_name", "vcf_path", "phenotype_codes"], + line.split("\t"))) + out.append("\t".join(row[c] for c in header)) + combined.write_text("\n".join(out) + "\n") + + cohort = oracle.Cohort(combined, bed_dir=data_dir / "beds") + for name in ("S10", "S11"): + assert name in cohort.samples + + _assert_matches(db, cohort, label="after add") + + +def test_after_add_samples_new_phenotype_matches_oracle(grown_db, data_dir, tmp_path): + """Restricted to the added samples, the numbers must describe only them.""" + db, grow_manifest = grown_db + cohort = oracle.Cohort(grow_manifest, bed_dir=data_dir / "beds") + _assert_matches(db, cohort, samples=["S10", "S11"], + phenotype=["ZZNEW"], label="added-only") diff --git a/tests/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000..d9f7c41 --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,252 @@ +"""Tests for variant Parquet layout resolution.""" +import pytest + +from afquery import storage + + +def _touch(path): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"") + + +@pytest.fixture +def variants(tmp_path): + d = tmp_path / "variants" + d.mkdir() + return d + + +# --------------------------------------------------------------------------- +# bucket_id +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("pos,expected", [ + (0, 0), + (1, 0), + (999_999, 0), + (1_000_000, 1), + (1_000_001, 1), + (2_500_000, 2), + (887_801, 0), +]) +def test_bucket_id_boundaries(pos, expected): + assert storage.bucket_id(pos) == expected + + +# --------------------------------------------------------------------------- +# chrom_layout +# --------------------------------------------------------------------------- + +def test_chrom_layout_partitioned(variants): + _touch(variants / "chr1" / "bucket_0.parquet") + assert storage.chrom_layout(variants, "chr1") == storage.PARTITIONED + + +def test_chrom_layout_flat(variants): + _touch(variants / "chr1.parquet") + assert storage.chrom_layout(variants, "chr1") == storage.FLAT + + +def test_chrom_layout_absent_uses_default(variants): + assert storage.chrom_layout(variants, "chr9", storage.FLAT) == storage.FLAT + assert storage.chrom_layout(variants, "chr9", storage.PARTITIONED) == storage.PARTITIONED + + +def test_chrom_layout_defaults_to_partitioned(variants): + assert storage.chrom_layout(variants, "chr9") == storage.PARTITIONED + + +def test_chrom_layout_bucket_directory_wins_over_flat(variants): + """The readers prefer the bucket directory, so the writer must too.""" + _touch(variants / "chr1.parquet") + _touch(variants / "chr1" / "bucket_0.parquet") + assert storage.chrom_layout(variants, "chr1") == storage.PARTITIONED + + +# --------------------------------------------------------------------------- +# detect_layout +# --------------------------------------------------------------------------- + +def test_detect_layout_empty_dir_is_partitioned(variants): + """A database with no variants yet must grow buckets, not a second layout.""" + assert storage.detect_layout(variants) == storage.PARTITIONED + + +def test_detect_layout_missing_dir_is_partitioned(tmp_path): + assert storage.detect_layout(tmp_path / "nope") == storage.PARTITIONED + + +def test_detect_layout_flat(variants): + _touch(variants / "chr1.parquet") + _touch(variants / "chr2.parquet") + assert storage.detect_layout(variants) == storage.FLAT + + +def test_detect_layout_partitioned(variants): + _touch(variants / "chr1" / "bucket_0.parquet") + assert storage.detect_layout(variants) == storage.PARTITIONED + + +def test_detect_layout_ignores_empty_chrom_dir(variants): + """A stray empty directory must not flip the database-wide answer.""" + (variants / "chr1").mkdir() + _touch(variants / "chr2.parquet") + assert storage.detect_layout(variants) == storage.FLAT + + +# --------------------------------------------------------------------------- +# non-canonical contig names +# --------------------------------------------------------------------------- + +def test_partitioned_chroms_keeps_non_canonical_contigs(variants): + """Unplaced and alt contigs are real data; filtering them here would hide it.""" + _touch(variants / "chr1" / "bucket_0.parquet") + _touch(variants / "chrGL000209.1" / "bucket_0.parquet") + _touch(variants / "chr1_KI270706v1_random" / "bucket_0.parquet") + assert storage.partitioned_chroms(variants) == { + "chr1", "chrGL000209.1", "chr1_KI270706v1_random", + } + + +def test_non_canonical_contig_resolves(variants): + _touch(variants / "chrGL000209.1" / "bucket_0.parquet") + assert storage.chrom_layout(variants, "chrGL000209.1") == storage.PARTITIONED + assert storage.variant_parquet_glob(variants, "chrGL000209.1") is not None + + +def test_flat_chroms(variants): + _touch(variants / "chr1.parquet") + _touch(variants / "chrX.parquet") + (variants / "chr2").mkdir() + assert storage.flat_chroms(variants) == {"chr1", "chrX"} + + +def test_existing_bucket_ids_handles_glob_metacharacters(variants): + """A contig name is a directory, never a pattern. + + GRCh38 spells HLA contigs with a literal asterisk. Interpolating one into a + glob makes it match its neighbours too, so the caller would merge and + recompute buckets belonging to a different contig. + """ + contig = "HLA-A*01:01" + _touch(variants / contig / "bucket_0.parquet") + _touch(variants / "HLA-A-99-01:01" / "bucket_5.parquet") + + assert storage.existing_bucket_ids(variants, contig) == [0] + assert storage.existing_bucket_ids(variants, "HLA-A-99-01:01") == [5] + + +def test_existing_bucket_ids_unknown_chrom(variants): + assert storage.existing_bucket_ids(variants, "chrNope") == [] + + +def test_stored_chroms_spans_both_layouts(variants): + """The Phase 2 recompute walks this set, so it must miss no chromosome.""" + _touch(variants / "chr1" / "bucket_0.parquet") + _touch(variants / "chr2" / "bucket_3.parquet") + _touch(variants / "chrX.parquet") + assert storage.stored_chroms(variants) == {"chr1", "chr2", "chrX"} + + +def test_stored_chroms_empty_database(variants): + assert storage.stored_chroms(variants) == set() + + +# --------------------------------------------------------------------------- +# mixed_layout_chroms +# --------------------------------------------------------------------------- + +def test_mixed_layout_chroms_detects_both(variants): + _touch(variants / "chr1.parquet") + _touch(variants / "chr1" / "bucket_0.parquet") + _touch(variants / "chr2" / "bucket_0.parquet") + _touch(variants / "chrX.parquet") + assert storage.mixed_layout_chroms(variants) == ["chr1"] + + +def test_mixed_layout_chroms_clean_db(variants): + _touch(variants / "chr1" / "bucket_0.parquet") + _touch(variants / "chr2" / "bucket_0.parquet") + assert storage.mixed_layout_chroms(variants) == [] + + +def test_mixed_layout_chroms_missing_dir(tmp_path): + assert storage.mixed_layout_chroms(tmp_path / "nope") == [] + + +# --------------------------------------------------------------------------- +# existing_bucket_ids +# --------------------------------------------------------------------------- + +def test_existing_bucket_ids_sorted_numerically(variants): + for b in (0, 2, 10, 1): + _touch(variants / "chr1" / f"bucket_{b}.parquet") + assert storage.existing_bucket_ids(variants, "chr1") == [0, 1, 2, 10] + + +def test_existing_bucket_ids_empty_for_flat(variants): + _touch(variants / "chr1.parquet") + assert storage.existing_bucket_ids(variants, "chr1") == [] + + +# --------------------------------------------------------------------------- +# path resolution +# --------------------------------------------------------------------------- + +def test_variant_parquet_for_pos_partitioned(variants): + _touch(variants / "chr1" / "bucket_2.parquet") + got = storage.variant_parquet_for_pos(variants, "chr1", 2_500_000) + assert got == variants / "chr1" / "bucket_2.parquet" + + +def test_variant_parquet_for_pos_missing_bucket(variants): + _touch(variants / "chr1" / "bucket_0.parquet") + assert storage.variant_parquet_for_pos(variants, "chr1", 5_000_000) is None + + +def test_variant_parquet_for_pos_flat(variants): + _touch(variants / "chr1.parquet") + assert storage.variant_parquet_for_pos(variants, "chr1", 500) == variants / "chr1.parquet" + + +def test_variant_parquet_for_pos_absent(variants): + assert storage.variant_parquet_for_pos(variants, "chr1", 500) is None + + +def test_variant_parquet_glob_partitioned(variants): + _touch(variants / "chr1" / "bucket_0.parquet") + assert storage.variant_parquet_glob(variants, "chr1").endswith("chr1/bucket_*.parquet") + + +def test_variant_parquet_glob_empty_chrom_dir(variants): + (variants / "chr1").mkdir() + assert storage.variant_parquet_glob(variants, "chr1") is None + + +def test_variant_parquet_glob_flat(variants): + _touch(variants / "chr1.parquet") + assert storage.variant_parquet_glob(variants, "chr1").endswith("chr1.parquet") + + +# --------------------------------------------------------------------------- +# iter_variant_parquets +# --------------------------------------------------------------------------- + +def test_iter_variant_parquets_covers_both_layouts(variants): + _touch(variants / "chr1.parquet") + _touch(variants / "chr2" / "bucket_0.parquet") + _touch(variants / "chr2" / "bucket_1.parquet") + got = {p.name for p in storage.iter_variant_parquets(variants)} + assert got == {"chr1.parquet", "bucket_0.parquet", "bucket_1.parquet"} + + +def test_iter_variant_parquets_skips_non_bucket_files(variants): + _touch(variants / "chr1" / "bucket_0.parquet") + _touch(variants / "chr1" / "notes.txt") + _touch(variants / "chr1" / "bucket_0.parquet.tmp") + got = [p.name for p in storage.iter_variant_parquets(variants)] + assert got == ["bucket_0.parquet"] + + +def test_iter_variant_parquets_missing_dir(tmp_path): + assert list(storage.iter_variant_parquets(tmp_path / "nope")) == [] diff --git a/tests/test_update.py b/tests/test_update.py index 8aa2a42..4d3b5eb 100644 --- a/tests/test_update.py +++ b/tests/test_update.py @@ -394,3 +394,57 @@ def test_check_sample_count_mismatch(fresh_db): results = check_database(fresh_db) non_info = [r for r in results if r.severity in ("warning", "error")] assert any("sample_count" in r.message or "mismatch" in r.message for r in non_info) + + +# --------------------------------------------------------------------------- +# Mixed layout: a chromosome stored in both layouts at once +# --------------------------------------------------------------------------- + +def _split_chr1(db: str) -> None: + """Put chr1 in both layouts, the state an older add-samples used to leave.""" + variants = os.path.join(db, "variants") + os.makedirs(os.path.join(variants, "chr1"), exist_ok=True) + shutil.copy( + os.path.join(variants, "chr1.parquet"), + os.path.join(variants, "chr1", "bucket_0.parquet"), + ) + + +def test_check_database_reports_mixed_layout(fresh_db): + _split_chr1(fresh_db) + + errors = [r for r in check_database(fresh_db) if r.severity == "error"] + assert any("chr1.parquet" in r.message for r in errors) + + +def test_add_samples_refuses_mixed_layout(fresh_db, tmp_path): + _split_chr1(fresh_db) + + vcf = str(tmp_path / "S10.vcf") + write_vcf(vcf, "S10", [("chr1", 7000, "A", "T", "0/1")]) + manifest = str(tmp_path / "manifest.tsv") + write_manifest(manifest, [("S10", "male", "wgs", vcf, "E11.9")]) + + with pytest.raises(UpdateError, match="Mixed variant layout"): + add_samples(fresh_db, manifest, threads=1) + + +def test_add_samples_manifest_sample_count_after_removal(fresh_db, tmp_path): + """sample_count is a count, not the next free id — they diverge after a removal.""" + remove_samples(fresh_db, ["S00"]) + + vcf = str(tmp_path / "S10.vcf") + write_vcf(vcf, "S10", [("chr1", 7000, "A", "T", "0/1")]) + manifest = str(tmp_path / "manifest.tsv") + write_manifest(manifest, [("S10", "male", "wgs", vcf, "E11.9")]) + add_samples(fresh_db, manifest, threads=1) + + con = sqlite3.connect(os.path.join(fresh_db, "metadata.sqlite")) + real_count = con.execute("SELECT COUNT(*) FROM samples").fetchone()[0] + con.close() + + data = json.loads(open(os.path.join(fresh_db, "manifest.json")).read()) + assert data["sample_count"] == real_count == 10 + assert data["next_sample_id"] == 11 + + assert [r for r in check_database(fresh_db) if r.severity == "error"] == [] diff --git a/tests/test_update_partitioned.py b/tests/test_update_partitioned.py new file mode 100644 index 0000000..023042c --- /dev/null +++ b/tests/test_update_partitioned.py @@ -0,0 +1,355 @@ +"""add-samples against a bucketed database. + +Every other add_samples test runs on the hand-built flat fixture in conftest.py, +which is the one layout create-db never produces. These build a real database +through run_preprocess, so the merge is exercised against the layout it will +actually meet. + +Cohort (tests/data): S00-S03 wgs, S04-S06 wes_kit_a (chr1:1000-2000, +chrX:1000-2000), S07-S09 wes_kit_b (chr1:3000-4000). chr1:1500 A>T is carried +by S00 het, S02 hom and S05 het. +""" +import shutil +import sqlite3 + +import pyarrow.parquet as pq +import pytest + +from afquery import storage +from afquery.bitmaps import deserialize +from afquery.database import Database +from afquery.preprocess import run_preprocess +from afquery.preprocess.update import add_samples, check_database + +from test_update import write_manifest, write_vcf + + +# --------------------------------------------------------------------------- +# Fixtures — build once per module, copy per test +# --------------------------------------------------------------------------- + +@pytest.fixture(scope="module") +def partitioned_src(tmp_path_factory, data_dir): + db = tmp_path_factory.mktemp("partitioned_src") + run_preprocess( + manifest_path=str(data_dir / "manifest.tsv"), output_dir=str(db), + genome_build="GRCh37", bed_dir=str(data_dir / "beds"), threads=2, + ) + return str(db) + + +@pytest.fixture +def partitioned_db(partitioned_src, tmp_path): + dest = tmp_path / "db" + shutil.copytree(partitioned_src, dest) + return str(dest) + + +@pytest.fixture(scope="module") +def covered_src(tmp_path_factory, data_dir): + """Same cohort with a coverage-evidence threshold, so filtered_bitmap is live.""" + db = tmp_path_factory.mktemp("covered_src") + run_preprocess( + manifest_path=str(data_dir / "manifest.tsv"), output_dir=str(db), + genome_build="GRCh37", bed_dir=str(data_dir / "beds"), threads=2, + min_covered=2, + ) + return str(db) + + +@pytest.fixture +def covered_db(covered_src, tmp_path): + dest = tmp_path / "covered" + shutil.copytree(covered_src, dest) + return str(dest) + + +def _add_one(db, tmp_path, name, variants, tech="wgs", sex="male", + phenotype="ZZNEW", bed_dir=None): + vcf = str(tmp_path / f"{name}.vcf") + write_vcf(vcf, name, variants) + manifest = str(tmp_path / f"{name}.tsv") + write_manifest(manifest, [(name, sex, tech, vcf, phenotype)]) + return add_samples(db, manifest, threads=1, bed_dir=bed_dir) + + +def _row(db, chrom, bucket, pos, ref, alt): + """One row of a bucket file as deserialized bitmaps, or None.""" + path = storage.bucket_path(f"{db}/variants", chrom, bucket) + if not path.exists(): + return None + t = pq.read_table(str(path)) + for i in range(len(t)): + if (t["pos"][i].as_py(), t["ref"][i].as_py(), t["alt"][i].as_py()) == (pos, ref, alt): + return { + name: deserialize(t[name][i].as_py()) + for name in ("het_bitmap", "hom_bitmap", "fail_bitmap", + "filtered_bitmap", "quality_pass_bitmap") + } + return None + + +# --------------------------------------------------------------------------- +# The regression: added samples must be visible to queries +# --------------------------------------------------------------------------- + +def test_add_samples_partitioned_new_variant_is_queryable(partitioned_db, tmp_path): + _add_one(partitioned_db, tmp_path, "S10", [("chr1", 7000, "A", "T", "0/1")]) + + db = Database(partitioned_db) + results = db.query(chrom="chr1", pos=7000, sex="both") + assert results, "variant added by add-samples is not visible to the query engine" + assert any(r.AC >= 1 for r in results) + + +def test_add_samples_partitioned_existing_variant_gains_carrier(partitioned_db, tmp_path): + before = Database(partitioned_db).query(chrom="chr1", pos=1500, sex="both") + ac_before = before[0].AC + an_before = before[0].AN + + # S10 is hom-alt, so it contributes 2 to AC and 2 to AN + _add_one(partitioned_db, tmp_path, "S10", [("chr1", 1500, "A", "T", "1/1")]) + + after = Database(partitioned_db).query(chrom="chr1", pos=1500, sex="both") + assert after[0].AC == ac_before + 2 + assert after[0].AN == an_before + 2 + + +def test_add_samples_partitioned_sets_carrier_bit(partitioned_db, tmp_path): + """The new sample id must be a set bit in the bucket, not merely counted.""" + _add_one(partitioned_db, tmp_path, "S10", [("chr1", 1500, "A", "T", "1/1")]) + + con = sqlite3.connect(f"{partitioned_db}/metadata.sqlite") + sid = con.execute("SELECT sample_id FROM samples WHERE sample_name='S10'").fetchone()[0] + con.close() + + row = _row(partitioned_db, "chr1", 0, 1500, "A", "T") + assert row is not None + assert sid in row["hom_bitmap"] + assert sid not in row["het_bitmap"] + + +def test_add_samples_partitioned_phenotype_isolated(partitioned_db, tmp_path): + """Filtered to the new sample alone the counts must describe that sample.""" + _add_one(partitioned_db, tmp_path, "S10", [("chr1", 1500, "A", "T", "1/1")]) + + res = Database(partitioned_db).query( + chrom="chr1", pos=1500, phenotype=["ZZNEW"], sex="both" + ) + assert len(res) == 1 + r = res[0] + assert (r.AC, r.AN, r.N_HET, r.N_HOM_ALT, r.N_HOM_REF) == (2, 2, 0, 1, 0) + + +# --------------------------------------------------------------------------- +# Layout is preserved, and never split +# --------------------------------------------------------------------------- + +def test_add_samples_preserves_chrom_layouts(partitioned_db, tmp_path): + variants = f"{partitioned_db}/variants" + before = {c: storage.chrom_layout(variants, c) + for c in storage.partitioned_chroms(variants) | storage.flat_chroms(variants)} + + _add_one(partitioned_db, tmp_path, "S10", [("chr1", 7000, "A", "T", "0/1")]) + + after = {c: storage.chrom_layout(variants, c) for c in before} + assert after == before + assert storage.mixed_layout_chroms(variants) == [] + + +def test_add_samples_partitioned_creates_new_bucket(partitioned_db, tmp_path): + assert 2 not in storage.existing_bucket_ids(f"{partitioned_db}/variants", "chr1") + + _add_one(partitioned_db, tmp_path, "S10", [("chr1", 2_500_000, "C", "G", "0/1")]) + + assert 2 in storage.existing_bucket_ids(f"{partitioned_db}/variants", "chr1") + assert _row(partitioned_db, "chr1", 2, 2_500_000, "C", "G") is not None + res = Database(partitioned_db).query(chrom="chr1", pos=2_500_000, sex="both") + assert res and res[0].AC == 1 + + +def test_add_samples_partitioned_new_chromosome_uses_buckets(partitioned_db, tmp_path): + variants = f"{partitioned_db}/variants" + assert "chr2" not in storage.partitioned_chroms(variants) + + _add_one(partitioned_db, tmp_path, "S10", [("chr2", 1000, "G", "A", "0/1")]) + + assert storage.chrom_layout(variants, "chr2") == storage.PARTITIONED + assert storage.existing_bucket_ids(variants, "chr2") == [0] + assert not storage.flat_path(variants, "chr2").exists() + res = Database(partitioned_db).query(chrom="chr2", pos=1000, sex="both") + assert res and res[0].AC == 1 + + +def test_add_samples_partitioned_check_database_clean(partitioned_db, tmp_path): + _add_one(partitioned_db, tmp_path, "S10", + [("chr1", 1500, "A", "T", "1/1"), ("chr1", 2_500_000, "C", "G", "0/1")]) + + errors = [r for r in check_database(partitioned_db) if r.severity == "error"] + assert errors == [] + + +# --------------------------------------------------------------------------- +# Phase 2: filtered_bitmap depends on the cohort, not on which bucket changed +# --------------------------------------------------------------------------- + +def test_add_samples_partitioned_recomputes_filtered_bitmap_in_untouched_buckets( + covered_db, tmp_path, data_dir +): + """A sample whose only variant is in bucket 2 still changes bucket 0. + + filtered_bitmap is (tech_bm - carriers) for every WES tech short of + min_covered quality-passing carriers. Enlarging wes_kit_a therefore moves it + at chr1:1500, which lives in a bucket the new sample contributed nothing to. + """ + before = _row(covered_db, "chr1", 0, 1500, "A", "T") + assert before is not None + assert before["filtered_bitmap"], "fixture should have a live filtered_bitmap" + + _add_one(covered_db, tmp_path, "S10", [("chr1", 2_500_000, "C", "G", "0/1")], + tech="wes_kit_a", bed_dir=str(data_dir / "beds")) + + con = sqlite3.connect(f"{covered_db}/metadata.sqlite") + sid = con.execute("SELECT sample_id FROM samples WHERE sample_name='S10'").fetchone()[0] + con.close() + + after = _row(covered_db, "chr1", 0, 1500, "A", "T") + assert sid in after["filtered_bitmap"], ( + "bucket 0 kept a stale filtered_bitmap: the new wes_kit_a sample is a " + "non-carrier of an under-covered tech and must be marked" + ) + assert before["filtered_bitmap"] < after["filtered_bitmap"] + + +def test_add_samples_untouched_chrom_is_not_rewritten_without_phase2( + partitioned_db, tmp_path): + """With no coverage threshold, nothing off the batch's chromosomes can change. + + filtered_bitmap is the only value that depends on the cohort as a whole, so + when it is not in play a chromosome the batch never mentions must be left + exactly as it was. + """ + chrX = storage.bucket_path(f"{partitioned_db}/variants", "chrX", 5) + before = chrX.read_bytes() + + _add_one(partitioned_db, tmp_path, "S10", + [("chr1", 2_500_000, "C", "G", "0/1")]) + + assert chrX.read_bytes() == before + + +# --------------------------------------------------------------------------- +# Phase 2 across chromosomes the batch never mentions +# --------------------------------------------------------------------------- + +@pytest.fixture(scope="module") +def two_chrom_wes_src(tmp_path_factory): + """A WES cohort on two chromosomes, under-covered on the second one. + + Only K0 carries chr2:1500, so with min_covered=2 the kit has too little + quality evidence there and every non-carrier of the kit is marked as lacking + coverage rather than counted homozygous reference. + """ + work = tmp_path_factory.mktemp("two_chrom_wes") + beds = work / "beds" + beds.mkdir() + (beds / "kit.bed").write_text("chr1\t999\t2000\nchr2\t999\t2000\n") + + entries = [] + for name, variants in [ + ("K0", [("chr1", 1500, "A", "T", "0/1"), ("chr2", 1500, "G", "C", "0/1")]), + ("K1", [("chr1", 1500, "A", "T", "1/1")]), + ("K2", [("chr1", 1500, "A", "T", "0/1")]), + ]: + vcf = str(work / f"{name}.vcf") + write_vcf(vcf, name, variants) + entries.append((name, "male", "kit", vcf, "COHORT")) + + manifest = str(work / "cohort.tsv") + write_manifest(manifest, entries) + + db = work / "db" + run_preprocess(manifest_path=manifest, output_dir=str(db), genome_build="GRCh37", + bed_dir=str(beds), threads=1, min_covered=2) + return str(db), str(beds), work + + +@pytest.fixture +def two_chrom_wes_db(two_chrom_wes_src, tmp_path): + src, beds, _work = two_chrom_wes_src + dest = tmp_path / "twochrom" + shutil.copytree(src, dest) + return str(dest), beds + + +def test_add_samples_recomputes_coverage_on_untouched_chromosomes( + two_chrom_wes_db, tmp_path): + """A WES sample added on chr1 must not be counted hom-ref on chr2. + + Its capture BED reaches chr2:1500, where the kit is under-covered, so the + sample belongs in N_NO_COVERAGE. Scoping the filtered_bitmap recomputation + to the chromosomes the batch happened to carry left chr2 stale and the + sample silently counted as homozygous reference, biasing the frequency. + """ + db_dir, beds = two_chrom_wes_db + before = Database(db_dir).query(chrom="chr2", pos=1500, sex="both")[0] + assert (before.N_HOM_REF, before.N_NO_COVERAGE) == (0, 2) + + _add_one(db_dir, tmp_path, "K3", [("chr1", 1500, "A", "T", "0/1")], + tech="kit", bed_dir=beds) + + after = Database(db_dir).query(chrom="chr2", pos=1500, sex="both")[0] + assert after.N_HOM_REF == 0, ( + "the added sample was counted homozygous reference on a chromosome the " + "batch never touched: its filtered_bitmap was left stale" + ) + assert after.N_NO_COVERAGE == 3 + + con = sqlite3.connect(f"{db_dir}/metadata.sqlite") + sid = con.execute("SELECT sample_id FROM samples WHERE sample_name='K3'").fetchone()[0] + con.close() + assert sid in _row(db_dir, "chr2", 0, 1500, "G", "C")["filtered_bitmap"] + + +def _merged_chroms(monkeypatch): + """Record which chromosomes add_samples asks the merge to visit.""" + from afquery.preprocess import update as update_mod + + seen = [] + original = update_mod._merge_chromosome_parquet + + def spy(chrom, *args, **kwargs): + seen.append(chrom) + return original(chrom, *args, **kwargs) + + monkeypatch.setattr(update_mod, "_merge_chromosome_parquet", spy) + return seen + + +def test_add_samples_wgs_batch_visits_only_its_own_chromosomes( + two_chrom_wes_db, tmp_path, monkeypatch): + """Only a batch that enlarges a capture tech can move coverage elsewhere. + + A WGS sample grows no tech bitmap, so nothing off its own chromosomes can + change and the rest of the store must not be read. The files would come out + byte-identical either way — the dirty guard sees to that — so what is + asserted here is that the work is not done at all. + """ + db_dir, _beds = two_chrom_wes_db + seen = _merged_chroms(monkeypatch) + + _add_one(db_dir, tmp_path, "W0", [("chr1", 1500, "A", "T", "0/1")]) + + assert seen == ["chr1"] + + +def test_add_samples_wes_batch_visits_the_whole_store( + two_chrom_wes_db, tmp_path, monkeypatch): + """A new capture sample moves coverage everywhere, so every chromosome is visited.""" + db_dir, beds = two_chrom_wes_db + seen = _merged_chroms(monkeypatch) + + _add_one(db_dir, tmp_path, "K3", [("chr1", 1500, "A", "T", "0/1")], + tech="kit", bed_dir=beds) + + assert sorted(seen) == ["chr1", "chr2"]