From ebbc178e3db510fa2b9421a59f4d8dcb51f57b63 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:04:48 +0200 Subject: [PATCH 01/14] refactor(storage): add a resolver for the variant Parquet layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rule that a chromosome's bucket directory takes priority over its flat Parquet file was reimplemented in every module that opens a variant file: the query engine, the dump and annotate workers, the benchmark helper, the compactor and the updater. Eight copies, none of them shared, and no single place to look when a writer and a reader disagree about where a row belongs. This adds the resolver as a standalone module. Nothing is migrated onto it yet, so the change is a pure addition; the call sites move over separately. partitioned_chroms() deliberately accepts any subdirectory rather than validating names against the canonical chromosome list. A database may hold unplaced or alt contigs — normalize_chrom() preserves bodies like 'chrGL000209.1' and 'chr1_KI270706v1_random' precisely so they stay outside the canonical set — and filtering them in the read path would make their data unreadable rather than merely untidy. --- src/afquery/storage.py | 166 +++++++++++++++++++++++++++++++ tests/test_storage.py | 221 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 387 insertions(+) create mode 100644 src/afquery/storage.py create mode 100644 tests/test_storage.py diff --git a/src/afquery/storage.py b/src/afquery/storage.py new file mode 100644 index 0000000..1f8bf63 --- /dev/null +++ b/src/afquery/storage.py @@ -0,0 +1,166 @@ +"""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 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.""" + ids: list[int] = [] + for p in Path(variants_dir).glob(f"{chrom}/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/test_storage.py b/tests/test_storage.py new file mode 100644 index 0000000..3921ac0 --- /dev/null +++ b/tests/test_storage.py @@ -0,0 +1,221 @@ +"""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"} + + +# --------------------------------------------------------------------------- +# 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")) == [] From 0adef0481eb38ec0f23b97c134878c4067e6596a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:07:12 +0200 Subject: [PATCH 02/14] fix(update): merge added samples into the partitioned bucket files _merge_chromosome_parquet only ever knew the flat layout: it read and wrote variants/.parquet and had no branch for the bucketed variants//bucket_N.parquet files. But create-db has produced the bucketed layout since it became the build default, so on every real database the merge found nothing to merge and wrote the incoming samples to a brand-new flat file instead. That file is dead weight. The readers resolve a chromosome's bucket directory first and never fall through to the flat file, so the added samples were never counted as carriers. They still counted toward AN through the capture index, which made them look homozygous reference at every position and pulled allele frequencies down across the whole database. Nothing failed, and afquery check had nothing to report. The read, merge, Phase 2 recompute, sort and write now live in one helper that operates on a single output file, and the chromosome-level function dispatches it either once for a flat file or once per affected bucket. The flat path is therefore exercised by the same code as the bucketed one. filtered_bitmap needs care: it is a function of the whole cohort rather than of which bucket received rows, because enlarging a WES technology changes (tech_bm - carrier_set) at every row of the chromosome. When min_covered is active, every existing bucket of an affected chromosome is revisited, and a per-row comparison keeps buckets whose bitmaps did not move from being rewritten. A chromosome new to the database follows the layout the database already uses, resolved once per update so that all chromosomes in one batch agree. --- src/afquery/preprocess/update.py | 209 ++++++++++++++++++++----------- 1 file changed, 136 insertions(+), 73 deletions(-) diff --git a/src/afquery/preprocess/update.py b/src/afquery/preprocess/update.py index 1a8144b..0de79c7 100644 --- a/src/afquery/preprocess/update.py +++ b/src/afquery/preprocess/update.py @@ -14,10 +14,11 @@ 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 +118,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 +162,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: @@ -213,8 +296,7 @@ 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 + prepared: list[tuple[int, str, str, BitMap, BitMap, BitMap, BitMap]] = [] for row in rows: if new_has_quality: @@ -225,17 +307,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 +323,43 @@ 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 has_phase2 and wes_tech_bitmaps and min_covered > 0: + # filtered_bitmap is a function of the whole cohort, not of which bucket + # received new rows: enlarging a WES tech changes (tech_bm - carrier_set) + # at every row of the chromosome. Buckets with no new data must still be + # revisited, or their N_NO_COVERAGE goes stale. + 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: @@ -710,12 +769,16 @@ def add_samples( if tech_obj is not None and tech_obj.bed_path is not None: wes_tech_bitmaps.setdefault(s.tech_id, BitMap()).add(s.sample_id) - # 10. Merge Parquet files + # 10. 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. + layout = storage.detect_layout(os.path.join(db_dir, "variants")) 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 From b3acce5aa60eb16f2ffb8c3203aa3f73d2d13c64 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:22:35 +0200 Subject: [PATCH 03/14] refactor(storage): route variant Parquet reads through the resolver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Eight call sites resolved a chromosome's storage layout for themselves: the query engine's path and glob helpers, the dump and annotate bucket workers, dump's chromosome availability check and work-unit enumeration, the benchmark variant sampler, the compactor, remove-samples, check-database, and the build resume check. Each reimplemented "a bucket directory wins over a flat file", and dump kept its own copy of BUCKET_SIZE. All of them now call the resolver. Behaviour is unchanged: the same precedence, the same per-chromosome decision, the same treatment of an empty chromosome directory. Verified by capturing dump, annotate, check, info, query, batch and region output for a bucketed and a flat database built from the same cohort, and diffing it against the same capture before the change — 22 files, all byte-identical. BUCKET_SIZE now has one definition, re-exported from build and dump so existing imports keep working. --- src/afquery/annotate.py | 7 +++--- src/afquery/benchmark.py | 35 +++++++++------------------- src/afquery/dump.py | 25 ++++++++------------ src/afquery/preprocess/build.py | 9 ++++---- src/afquery/preprocess/compact.py | 10 ++------ src/afquery/preprocess/update.py | 38 +++++++------------------------ src/afquery/query.py | 30 +++++------------------- 7 files changed, 44 insertions(+), 110 deletions(-) diff --git a/src/afquery/annotate.py b/src/afquery/annotate.py index 707e256..51cdd86 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 @@ -52,8 +53,8 @@ def _compute_chunk_annotations( bucket_end = (bucket_id + 1) * 1_000_000 - 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( 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..5ffc318 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)) diff --git a/src/afquery/preprocess/update.py b/src/afquery/preprocess/update.py index 0de79c7..b9d4dee 100644 --- a/src/afquery/preprocess/update.py +++ b/src/afquery/preprocess/update.py @@ -7,7 +7,6 @@ import time from dataclasses import dataclass from datetime import datetime, timezone -from pathlib import Path import duckdb import pyarrow as pa @@ -862,22 +861,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)) @@ -1013,20 +1001,10 @@ 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"))) - ) - - 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..fa611dc 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) @@ -312,23 +304,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 From 18be233e70e195e24a954bfc9d839604b7a561af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:24:00 +0200 Subject: [PATCH 04/14] fix(check): report chromosomes stored in two layouts at once MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A database updated by the previous add-samples ends up holding both variants/chr1/ and variants/chr1.parquet, and nothing said so. check-database validated each file on its own merits, and each is individually well-formed; the damage is in the pair, because queries read only the bucket directory and every sample whose calls live in the flat file is counted as homozygous reference. check now errors once per affected chromosome and names that consequence, so the failure is legible without opening the Parquet files. add-samples refuses to run on such a database, and does so before ingest starts rather than after an operator has waited out a long VCF pass. remove-samples is deliberately left unguarded. It already clears bits from both layouts, is safe on a mixed database, and is the first step of the recovery procedure — refusing there would leave no way out. --- src/afquery/preprocess/update.py | 48 ++++++++++++++++++++++++-------- 1 file changed, 36 insertions(+), 12 deletions(-) diff --git a/src/afquery/preprocess/update.py b/src/afquery/preprocess/update.py index b9d4dee..ae41015 100644 --- a/src/afquery/preprocess/update.py +++ b/src/afquery/preprocess/update.py @@ -659,12 +659,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) @@ -672,7 +686,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() } @@ -683,7 +697,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") @@ -699,7 +713,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) @@ -735,7 +749,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_") @@ -743,7 +757,7 @@ 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) @@ -768,7 +782,7 @@ def add_samples( if tech_obj is not None and tech_obj.bed_path is not None: wes_tech_bitmaps.setdefault(s.tech_id, BitMap()).add(s.sample_id) - # 10. Merge Parquet files. The layout is decided once for the whole + # 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. layout = storage.detect_layout(os.path.join(db_dir, "variants")) @@ -787,7 +801,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)" @@ -802,10 +816,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( @@ -817,7 +831,7 @@ def add_samples( 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") @@ -1001,6 +1015,16 @@ def info(msg: str) -> None: err("variants/ directory not found") return results + # 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 = [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) From fec399cf790bbdf961fd44245977c5db3f5073cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:25:00 +0200 Subject: [PATCH 05/14] fix(update): write the real sample count to the manifest after add MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit add-samples passed the next free sample id where _update_manifest expects the sample count. The two agree only until a sample is removed, because ids are never reused: after a removal the next id runs ahead of the count. remove-samples already writes SELECT COUNT(*), so the mismatch only appears on a remove-then-add cycle — which is exactly what the recovery procedure for a split variant layout asks operators to perform. Without this, following that procedure ends with check-database warning about a sample_count mismatch the operator caused by doing as they were told. --- src/afquery/preprocess/update.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/afquery/preprocess/update.py b/src/afquery/preprocess/update.py index ae41015..dc7aeaf 100644 --- a/src/afquery/preprocess/update.py +++ b/src/afquery/preprocess/update.py @@ -828,6 +828,10 @@ 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() @@ -836,7 +840,7 @@ def add_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)) From 88730534c8ef294afee96db72bc381b25fb248e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:28:46 +0200 Subject: [PATCH 06/14] test(update): cover add-samples against a partitioned database MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every existing add_samples test copies the hand-built fixture from conftest, which writes a flat variants/.parquet directly. That is the one layout create-db does not produce, so a merge that understood only the flat layout passed the whole suite while corrupting every real database it touched. These build a database through run_preprocess and assert, through the query engine, that an added sample is counted as a carrier — for a new variant, for a variant the cohort already had, and at bit level in the bucket file, which no test checked after an add before. Also covered: a position past the last existing bucket, a chromosome absent from the database, and that the per chromosome layout is unchanged by the add and never split in two. The coverage-evidence case adds a wes_kit_a sample whose only variant lives in bucket 2 and asserts that bucket 0's filtered_bitmap picked it up, since that bitmap is a function of the cohort rather than of which bucket received rows. A partial fix that only rewrote buckets with new data would pass everything else and fail that one. Eight of the ten fail against the previous merge; the two that do not are guards against a fix that rewrites more than it should. --- tests/test_update.py | 54 +++++++ tests/test_update_partitioned.py | 232 +++++++++++++++++++++++++++++++ 2 files changed, 286 insertions(+) create mode 100644 tests/test_update_partitioned.py 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..83543bf --- /dev/null +++ b/tests/test_update_partitioned.py @@ -0,0 +1,232 @@ +"""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_partitioned_untouched_chrom_is_not_rewritten(covered_db, tmp_path, data_dir): + """Chromosomes the batch never touches keep their files as they were.""" + chrX = storage.bucket_path(f"{covered_db}/variants", "chrX", 5) + before = chrX.read_bytes() + + _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")) + + assert chrX.read_bytes() == before From 8f3bf8ebd0ba7846475b6c68331dfac7a2ea1f2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:37:31 +0200 Subject: [PATCH 07/14] fix(query): drop samples with no allele from the eligible set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On chrY the eligible set included females, so they were reported as homozygous reference and counted in n_samples_eligible. A female has no chrY to genotype: she is neither a carrier nor homozygous reference there. AN has always excluded such samples, which left the two disagreeing — in the test cohort a chrY variant reported AN=1 alongside n_eligible=2, and the per-genotype counts summed to more samples than could hold an allele. Fixed where eligibility is decided rather than in the five places that derive N_HOM_REF from it, so point, batch, region, dump, annotate and variant-info all pick it up. split_ploidy already knows which samples carry alleles at a position; intersecting with it leaves autosomes, chrX and chrM untouched, since there both sexes are genotyped and only the ploidy differs. The golden case for chrY recorded the old n_eligible and is updated. Found by the query oracle added in the following commit. --- src/afquery/query.py | 13 +++++++++++++ tests/data/expected_results.json | 4 ++-- tests/test_haploid_stats.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/src/afquery/query.py b/src/afquery/query.py index fa611dc..4939c30 100644 --- a/src/afquery/query.py +++ b/src/afquery/query.py @@ -294,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, 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/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 From 3c77591b321b35972a96332c338ccbabb8ccefe4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:41:41 +0200 Subject: [PATCH 08/14] fix(compact): survive a Parquet file whose rows are all dropped MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit compact selected the rows to keep with table.take(keep_indices) on a plain Python list. When every row of a file is dropped the list is empty, pyarrow infers a null-typed index array, and take has no kernel for (uint32, null) — so the command died with ArrowNotImplementedError partway through, leaving the database half compacted. Emptying a whole file is ordinary: remove a sample that was the only carrier of everything in a 1 Mbp bucket and that bucket has nothing left to keep. Spelling out the index type makes the empty case produce an empty table, which is what the rest of the function already handles. --- src/afquery/preprocess/compact.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/afquery/preprocess/compact.py b/src/afquery/preprocess/compact.py index 5ffc318..618ec59 100644 --- a/src/afquery/preprocess/compact.py +++ b/src/afquery/preprocess/compact.py @@ -109,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"], From 1a1c97df7f7fde9b79d230fcc3c782302dc23c18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:41:51 +0200 Subject: [PATCH 09/14] test: check query results against an independent oracle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The suite was green while add-samples silently dropped every added sample, because it only ever asked whether the code agreed with a fixture the same project had written by hand. Nothing compared the numbers against the VCF, BED and manifest text the database was built from. tests/oracle.py works those numbers out from the raw files and imports nothing from afquery — reusing the code under test to predict its own output would only restate its assumptions. The coverage half is the half that matters: AC and the genotype counts come from the VCFs, but AN, N_HOM_REF and N_NO_COVERAGE come from capture-BED intervals crossed with sex-dependent ploidy, which is where every silent counting bug this project has had actually lived. It is checked against a bucketed build, a flat build, a phenotype subset, and a database grown by add-samples. tests/test_invariants.py adds properties that compare one database against another rather than against a recorded expectation: building a cohort in one pass equals building half and adding the rest; the storage layout is not observable through any query; add-then-remove returns the answers to what they were; and compact changes nothing that still has a carrier, twice over. Between them they found the chrY eligibility bug and the compact crash fixed in the two preceding commits. --- tests/oracle.py | 195 +++++++++++++++++++++++++++++++ tests/test_invariants.py | 186 +++++++++++++++++++++++++++++ tests/test_oracle_consistency.py | 164 ++++++++++++++++++++++++++ 3 files changed, 545 insertions(+) create mode 100644 tests/oracle.py create mode 100644 tests/test_invariants.py create mode 100644 tests/test_oracle_consistency.py 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_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") From 0dbfb6936b2f3d0842fd200d258f2a0d11b9fe6e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Tue, 8 Sep 2026 15:42:50 +0200 Subject: [PATCH 10/14] docs: document the variant layouts and the mixed-layout recovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The flat layout was undocumented, which is part of why a write path that only understood it survived: the data model described bucketed storage as if it were the only kind. It now describes both, says which one create-db produces, and records that a chromosome must never carry two. Troubleshooting gains the entry for databases already split by the previous add-samples: how to recognise it — AN grows, AC does not — and the remove, delete-orphans, re-add sequence that repairs it without a full rebuild. The order matters, so the recovery records the phenotype rows before removal deletes them. The ploidy page already gave chrY a female ploidy of 0; it now says what that means for eligibility, which is that n_eligible and N_HOM_REF count males only there. --- docs/advanced/ploidy-and-sex-chroms.md | 2 + docs/guides/update-database.md | 9 ++-- docs/reference/data-model.md | 16 ++++++ docs/troubleshooting.md | 72 ++++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 3 deletions(-) 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..e73a6ef 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,10 @@ 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`. The recomputation runs across every bucket of the +chromosomes touched by the new samples, because the value depends on the whole +cohort rather than on which bucket received rows; chromosomes the batch does not +touch are not rewritten. 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..abdf7dd 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -81,6 +81,77 @@ 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 orphan files before deleting anything. +find ./db/variants -maxdepth 1 -name '*.parquet' + +# 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 the orphan flat files. Each one has a sibling directory of the same name. +find ./db/variants -maxdepth 1 -name '*.parquet' -delete + +# 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 +187,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) | --- From b4d99c2824ea3ea0b081a2e029041e87c6720463 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Wed, 9 Sep 2026 09:15:48 +0200 Subject: [PATCH 11/14] fix(annotate): take the bucket size from the layout resolver annotate.py computed bucket boundaries from a hardcoded 1_000_000 in three places while resolving the file path through storage. If the bucket size ever changed, storage.bucket_path would open one bucket and the WHERE pos BETWEEN range would filter another, so annotation would report AN=0 for most variants with no error. The resolver exists so readers and writers cannot drift; these were the last sites still holding their own copy of the rule. --- src/afquery/annotate.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/afquery/annotate.py b/src/afquery/annotate.py index 51cdd86..48e501f 100644 --- a/src/afquery/annotate.py +++ b/src/afquery/annotate.py @@ -49,8 +49,8 @@ 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 storage.chrom_layout(_db / "variants", chrom) == storage.PARTITIONED: @@ -186,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) From f93348ff7da05a78c32fed6fc139de3503c1bc6f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Wed, 9 Sep 2026 09:15:48 +0200 Subject: [PATCH 12/14] docs: keep flat variant files that have no bucket directory The recovery procedure deleted every top-level Parquet file with a blanket find -delete, on the assumption that each one has a sibling bucket directory. A broken update-db that introduced a chromosome absent from the original build wrote a flat file with no sibling: check does not flag it because the layout is not mixed, queries read it correctly through the flat fallback, and it is the only copy of that chromosome. Deleting it would drop the chromosome silently. Both the listing and the delete now test for the sibling directory. --- docs/troubleshooting.md | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index abdf7dd..7ff8ca8 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -120,16 +120,23 @@ sqlite3 ./db/metadata.sqlite \ 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 orphan files before deleting anything. -find ./db/variants -maxdepth 1 -name '*.parquet' +# 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 the orphan flat files. Each one has a sibling directory of the same name. -find ./db/variants -maxdepth 1 -name '*.parquet' -delete +# 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/ From 972954bb2f9d60d06867dddf023dfd23344fa2a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Wed, 9 Sep 2026 09:23:00 +0200 Subject: [PATCH 13/14] fix(update): recompute coverage evidence across the whole database MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit add-samples recomputed filtered_bitmap only on the chromosomes the batch carried variants on. But filtered_bitmap is (tech_bm - carriers) row by row, so enlarging a WES tech moves it at every row of every chromosome, not only where the new rows landed. A sample added with variants on chr1 was therefore left out of the coverage bitmap on chr2, and counted as homozygous reference wherever its capture BED reached — silently, and biasing the frequency down. The answer depended on nothing but which chromosome the batch happened to mention. The merge now walks every chromosome the database holds whenever a WES tech grows, and _merge_chromosome_parquet accepts a chromosome with no new rows as a recompute-only pass. The per-file dirty guard already suppresses writes where no bitmap moved, so a WGS-only batch still rewrites nothing. A test asserted the old behaviour was intended; it now covers the case where leaving a chromosome alone is genuinely correct, which is when no coverage threshold is set. --- docs/guides/update-database.md | 12 +++-- src/afquery/preprocess/update.py | 31 +++++++---- src/afquery/storage.py | 5 ++ tests/test_storage.py | 12 +++++ tests/test_update_partitioned.py | 89 ++++++++++++++++++++++++++++++-- 5 files changed, 130 insertions(+), 19 deletions(-) diff --git a/docs/guides/update-database.md b/docs/guides/update-database.md index e73a6ef..c28fa7e 100644 --- a/docs/guides/update-database.md +++ b/docs/guides/update-database.md @@ -69,10 +69,14 @@ 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 across every bucket of the -chromosomes touched by the new samples, because the value depends on the whole -cohort rather than on which bucket received rows; chromosomes the batch does not -touch 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. Files +whose bitmaps do not actually move are left untouched, so adding WGS-only +samples costs a read and no rewrites, while 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/src/afquery/preprocess/update.py b/src/afquery/preprocess/update.py index dc7aeaf..dec3113 100644 --- a/src/afquery/preprocess/update.py +++ b/src/afquery/preprocess/update.py @@ -284,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) @@ -295,6 +291,15 @@ 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) + # 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: @@ -341,11 +346,9 @@ def _merge_chromosome_parquet( by_bucket.setdefault(storage.bucket_id(prepared_row[0]), []).append(prepared_row) target_ids = set(by_bucket) - if has_phase2 and wes_tech_bitmaps and min_covered > 0: - # filtered_bitmap is a function of the whole cohort, not of which bucket - # received new rows: enlarging a WES tech changes (tech_bm - carrier_set) - # at every row of the chromosome. Buckets with no new data must still be - # revisited, or their N_NO_COVERAGE goes stale. + 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)) total_new = 0 @@ -785,7 +788,15 @@ def add_samples( # 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. - layout = storage.detect_layout(os.path.join(db_dir, "variants")) + variants_dir = os.path.join(db_dir, "variants") + if wes_tech_bitmaps: + # 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. + 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, diff --git a/src/afquery/storage.py b/src/afquery/storage.py index 1f8bf63..0526c93 100644 --- a/src/afquery/storage.py +++ b/src/afquery/storage.py @@ -71,6 +71,11 @@ def flat_chroms(variants_dir: Path | str) -> set[str]: 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. diff --git a/tests/test_storage.py b/tests/test_storage.py index 3921ac0..a30189b 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -121,6 +121,18 @@ def test_flat_chroms(variants): assert storage.flat_chroms(variants) == {"chr1", "chrX"} +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 # --------------------------------------------------------------------------- diff --git a/tests/test_update_partitioned.py b/tests/test_update_partitioned.py index 83543bf..4e35c88 100644 --- a/tests/test_update_partitioned.py +++ b/tests/test_update_partitioned.py @@ -221,12 +221,91 @@ def test_add_samples_partitioned_recomputes_filtered_bitmap_in_untouched_buckets assert before["filtered_bitmap"] < after["filtered_bitmap"] -def test_add_samples_partitioned_untouched_chrom_is_not_rewritten(covered_db, tmp_path, data_dir): - """Chromosomes the batch never touches keep their files as they were.""" - chrX = storage.bucket_path(f"{covered_db}/variants", "chrX", 5) +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(covered_db, tmp_path, "S10", [("chr1", 2_500_000, "C", "G", "0/1")], - tech="wes_kit_a", bed_dir=str(data_dir / "beds")) + _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"] From d055493425ad7462d48c2b56446a3f9fca2e7b0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20L=C3=B3pez=20L=C3=B3pez?= Date: Wed, 9 Sep 2026 09:28:55 +0200 Subject: [PATCH 14/14] fix(update): visit the whole store only when a capture tech grows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The recompute across every chromosome keyed off wes_tech_bitmaps, which is built from every existing sample that has a BED. In any database holding one capture sample that made a WGS-only batch read and deserialize the entire store, when no bitmap off its own chromosomes could move: filtered_bitmap is (tech_bm - carriers), and on an untouched chromosome carriers cannot change, so only a batch that puts a sample into a capture technology moves anything. Gate on that instead. storage.existing_bucket_ids now treats the chromosome as a directory rather than interpolating it into a glob pattern. It was previously only ever reached with names filtered against the known set; it is now also reached with names read straight off disk, where a contig spelled with an asterisk — GRCh38 does this for HLA — would pull in the buckets of its neighbours. --- docs/guides/update-database.md | 9 ++++--- src/afquery/preprocess/update.py | 8 ++++-- src/afquery/storage.py | 8 +++++- tests/test_storage.py | 19 ++++++++++++++ tests/test_update_partitioned.py | 44 ++++++++++++++++++++++++++++++++ 5 files changed, 81 insertions(+), 7 deletions(-) diff --git a/docs/guides/update-database.md b/docs/guides/update-database.md index c28fa7e..35606f0 100644 --- a/docs/guides/update-database.md +++ b/docs/guides/update-database.md @@ -73,10 +73,11 @@ 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. Files -whose bitmaps do not actually move are left untouched, so adding WGS-only -samples costs a read and no rewrites, while adding a WES sample to a large -database rewrites broadly and takes minutes rather than seconds. +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/src/afquery/preprocess/update.py b/src/afquery/preprocess/update.py index dec3113..b36e49f 100644 --- a/src/afquery/preprocess/update.py +++ b/src/afquery/preprocess/update.py @@ -767,6 +767,7 @@ def add_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" @@ -783,18 +784,21 @@ 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) # 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: + 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. + # 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: diff --git a/src/afquery/storage.py b/src/afquery/storage.py index 0526c93..9a2b708 100644 --- a/src/afquery/storage.py +++ b/src/afquery/storage.py @@ -126,8 +126,14 @@ def mixed_layout_chroms(variants_dir: Path | str) -> list[str]: 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 Path(variants_dir).glob(f"{chrom}/bucket_*.parquet"): + for p in chrom_dir.glob("bucket_*.parquet"): stem = p.stem[len("bucket_"):] if stem.isdigit(): ids.append(int(stem)) diff --git a/tests/test_storage.py b/tests/test_storage.py index a30189b..d9f7c41 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -121,6 +121,25 @@ def test_flat_chroms(variants): 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") diff --git a/tests/test_update_partitioned.py b/tests/test_update_partitioned.py index 4e35c88..023042c 100644 --- a/tests/test_update_partitioned.py +++ b/tests/test_update_partitioned.py @@ -309,3 +309,47 @@ def test_add_samples_recomputes_coverage_on_untouched_chromosomes( 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"]