feat: support for bulk writing allele frequency records (#64)

Author: theferrit32

src/anyvlm/functions/ingest_vcf.py

@@ -31,7 +31,10 @@ class VcfAfColumnsError(Exception):
 def _yield_expression_af_batches(
     vcf: pysam.VariantFile, batch_size: int = 1000
 ) -> Iterator[list[tuple[str, AfData]]]:
-    """Generate a variant expression-allele frequency data pairing, one at a time
+    """Generate batches of tuples of (variant expression, allele frequency data).
+
+    Operates lazily so only one batch is in memory at a time. If a VCF record has
+    multiple alternate alleles, each is returned as a separate item.
 
     :param vcf: VCF to pull variants from
     :param batch_size: size of return batches
@@ -98,6 +101,8 @@ def ingest_vcf(
     for batch in _yield_expression_af_batches(vcf):
         expressions, afs = zip(*batch, strict=True)
         variant_ids = av.put_allele_expressions(expressions, assembly)
+
+        cafs = []
         for variant_id, af in zip(variant_ids, afs, strict=True):
             if variant_id is None:
                 continue
@@ -114,4 +119,6 @@ def ingest_vcf(
                 ),
                 cohort=StudyGroup(name="rare disease"),
             )
-            storage.add_allele_frequencies(caf)
+            cafs.append(caf)
+
+        storage.add_allele_frequencies(cafs)

src/anyvlm/storage/base_storage.py

@@ -30,14 +30,12 @@ def sanitized_url(self) -> str:
         """Return a sanitized URL (password masked) of the database connection string."""
 
     @abstractmethod
-    def add_allele_frequencies(self, caf: AnyVlmCohortAlleleFrequencyResult) -> None:
+    def add_allele_frequencies(
+        self, cafs: list[AnyVlmCohortAlleleFrequencyResult]
+    ) -> None:
         """Add allele frequency data to the database. Will skip conflicts.
 
-        NOTE: For now, this will only insert a single caf record into the database.
-        Single insertion is used to do a simple test of the storage backend.
-        Issue-34 will support batch insertion of caf records.
-
-        :param caf: Cohort allele frequency study result object to insert into the DB
+        :param cafs: List of cohort allele frequency study result objects to insert
         """
 
     @abstractmethod

src/anyvlm/storage/postgres.py

@@ -54,20 +54,21 @@ def sanitized_url(self) -> str:
             netloc += f":{parsed.port}"
         return f"{parsed.scheme}://{netloc}{parsed.path}"
 
-    def add_allele_frequencies(self, caf: AnyVlmCohortAlleleFrequencyResult) -> None:
+    def add_allele_frequencies(
+        self, cafs: list[AnyVlmCohortAlleleFrequencyResult]
+    ) -> None:
         """Add allele frequency data to the database. Will skip conflicts.
 
-        NOTE: For now, this will only insert a single caf record into the database.
-        Single insertion is used to do a simple test of the storage backend.
-        Issue-34 will support batch insertion of caf records.
-
-        :param caf: Cohort allele frequency study result object to insert into the DB
+        :param cafs: List of cohort allele frequency study result objects to insert
         """
-        db_entity = mapper_registry.to_db_entity(caf)
+        if not cafs:
+            return
+
+        db_entities = [mapper_registry.to_db_entity(caf) for caf in cafs]
         stmt = insert(orm.AlleleFrequencyData).on_conflict_do_nothing()
 
         with self.session_factory() as session, session.begin():
-            session.execute(stmt, db_entity.to_dict())
+            session.execute(stmt, [entity.to_dict() for entity in db_entities])
 
     def get_caf_by_vrs_allele_id(
         self, vrs_allele_id: str

tests/integration/functions/test_get_caf.py

@@ -75,9 +75,11 @@ def populated_postgres_storage(
     caf_iri: AnyVlmCohortAlleleFrequencyResult,
 ):
     """Populate the postgres storage with allele frequencies for testing"""
-    for allele in alleles.values():
-        caf = build_caf(caf_iri, allele_id=allele["variation"]["id"])
-        postgres_storage.add_allele_frequencies(caf)
+    cafs = [
+        build_caf(caf_iri, allele_id=allele["variation"]["id"])
+        for allele in alleles.values()
+    ]
+    postgres_storage.add_allele_frequencies(cafs)
     return postgres_storage
 
 

tests/integration/storage/test_postgres_integration.py

@@ -22,7 +22,7 @@ def test_db_lifecycle(
     caf_rows = return_cafs(storage)
     assert caf_rows == []
 
-    storage.add_allele_frequencies(caf_iri)
+    storage.add_allele_frequencies([caf_iri])
     caf_rows = return_cafs(storage)
     assert len(caf_rows) == 1
 

tests/unit/storage/test_postgres_unit.py

@@ -62,7 +62,7 @@ def test_add_allele_frequencies(
     """Test that add_allele_frequencies method works correctly"""
     caf = request.getfixturevalue(caf_fixture_name)
     try:
-        postgres_storage.add_allele_frequencies(caf)
+        postgres_storage.add_allele_frequencies([caf])
     except Exception as e:  # noqa: BLE001
         pytest.fail(f"add_allele_frequencies raised an exception: {e}")
 
@@ -75,7 +75,7 @@ def test_add_allele_frequencies(
         cohort=StudyGroup(name="rare disease"),  # type: ignore
     )  # type: ignore
 
-    postgres_storage.add_allele_frequencies(caf)
+    postgres_storage.add_allele_frequencies([caf])
 
 
 def test_add_allele_frequencies_failures(
@@ -84,4 +84,4 @@ def test_add_allele_frequencies_failures(
 ):
     """Test that add_allele_frequencies method fails correctly on bad input"""
     with pytest.raises(IntegrityError, match='null value in column "cohort"'):
-        postgres_storage.add_allele_frequencies(caf_empty_cohort)
+        postgres_storage.add_allele_frequencies([caf_empty_cohort])