Bugfixes for reference data update (#690)

* add logger

* ws

* lint

* quiet loggers

* pyproj

* Fix bug

* Handle enums correctly

* lint

* Improve logging

* another

* account for key in selects compare

* i guess this is fine for now?

* Leave a comment

* Question

* Bugfixes

* almost there

* another batch

* ugh

* lint

* Import enums

* Missed one

* support dropping a dataset

* lint

Author: bpblanken

v03_pipeline/lib/annotations/mito.py

@@ -44,7 +44,7 @@ def HL(mt: hl.MatrixTable, **_: Any) -> hl.Expression:  # noqa: N802
     return hl.if_else(is_called, mt.HL, 0)
 
 
-def high_constraint_region(
+def high_constraint_region_mito(
     ht: hl.Table,
     interval_ht: hl.Table,
     **_: Any,

v03_pipeline/lib/model/dataset_type.py

@@ -175,7 +175,7 @@ def formatting_annotation_fns(
             DatasetType.MITO: [
                 mito.common_low_heteroplasmy,
                 mito.haplogroup,
-                mito.high_constraint_region,
+                mito.high_constraint_region_mito,
                 mito.mitotip,
                 mito.rsid,
                 shared.sorted_transcript_consequences,

v03_pipeline/lib/model/reference_dataset_collection.py

@@ -54,7 +54,7 @@ def datasets(self, dataset_type: DatasetType) -> list[str]:
             (ReferenceDatasetCollection.INTERVAL, DatasetType.MITO): [
                 'high_constraint_region_mito',
             ],
-        }[(self, dataset_type)]
+        }.get((self, dataset_type), [])
 
     def table_key_type(
         self,

v03_pipeline/lib/reference_data/clinvar.py

@@ -1,4 +1,5 @@
 import gzip
+import os
 import subprocess
 import tempfile
 import urllib
@@ -7,6 +8,7 @@
 
 from v03_pipeline.lib.annotations.enums import CLINVAR_PATHOGENICITIES_LOOKUP
 from v03_pipeline.lib.logger import get_logger
+from v03_pipeline.lib.model import Env
 from v03_pipeline.lib.model.definitions import ReferenceGenome
 
 CLINVAR_ASSERTIONS = [
@@ -104,7 +106,7 @@ def download_and_import_latest_clinvar_vcf(
 
     with tempfile.NamedTemporaryFile(suffix='.vcf.gz', delete=False) as tmp_file:
         urllib.request.urlretrieve(clinvar_url, tmp_file.name)  # noqa: S310
-        gcs_tmp_file_name = f'gs://seqr-scratch-temp{tmp_file.name}'
+        gcs_tmp_file_name = os.path.join(Env.HAIL_TMPDIR, os.path.basename(tmp_file.name))
         safely_move_to_gcs(tmp_file.name, gcs_tmp_file_name)
         mt = hl.import_vcf(
             gcs_tmp_file_name,

v03_pipeline/lib/reference_data/compare_globals.py

@@ -11,6 +11,7 @@
 from v03_pipeline.lib.reference_data.config import CONFIG
 from v03_pipeline.lib.reference_data.dataset_table_operations import (
     get_all_select_fields,
+    get_enum_select_fields,
     get_ht_path,
     import_ht_from_config_path,
     parse_dataset_version,
@@ -50,9 +51,9 @@ def from_dataset_configs(
                 ),
             )
             enums[dataset] = dataset_config.get('enum_select', {})
-            selects[dataset] = set(
-                get_all_select_fields(dataset_ht, dataset_config).keys(),
-            )
+            dataset_ht = dataset_ht.select(**get_all_select_fields(dataset_ht, dataset_config))
+            dataset_ht = dataset_ht.transmute(**get_enum_select_fields(dataset_ht, dataset_config))
+            selects[dataset] = set(dataset_ht.row) - set(dataset_ht.key)
         return cls(paths, versions, enums, selects)
 
     @classmethod
@@ -65,12 +66,14 @@ def from_ht(
         rdc_globals_struct = hl.eval(ht.globals)
         paths = dict(rdc_globals_struct.paths)
         versions = dict(rdc_globals_struct.versions)
-        enums = dict(rdc_globals_struct.enums)
+        # enums are nested structs
+        enums = {k: dict(v) for k,v in rdc_globals_struct.enums.items()}
 
         selects = {}
         for dataset in rdc.datasets(dataset_type):
             if dataset in ht.row:
-                selects[dataset] = set(ht[dataset])
+                # NB: handle an edge case (mito high constraint) where we annotate a bool from the reference dataset collection
+                selects[dataset] = set(ht[dataset]) if isinstance(ht[dataset], hl.StructExpression) else set()
         return cls(paths, versions, enums, selects)
 
 

v03_pipeline/lib/reference_data/compare_globals_test.py

@@ -22,15 +22,15 @@ class CompareGlobalsTest(unittest.TestCase):
                 '38': {
                     'custom_import': None,
                     'source_path': 'a_path',  # 'a' has a custom import
-                    'select': {'test_select': 'info.test_select'},
+                    'select': {'test_select': 'info.test_select', 'test_enum': 'test_enum'},
                     'version': 'a_version',
                     'enum_select': {'test_enum': ['A', 'B']},
                 },
             },
             'b': {  # b is missing version
                 '38': {
                     'path': 'b_path',
-                    'select': {'test_select': 'info.test_select'},
+                    'select': {'test_select': 'info.test_select', 'test_enum': 'test_enum'},
                     'enum_select': {'test_enum': ['C', 'D']},
                     'custom_select': lambda ht: {'field2': ht.info.test_select_2},
                 },
@@ -56,8 +56,10 @@ def test_create_globals_from_dataset_ht_configs(
                     info=hl.tstruct(
                         test_select=hl.tint,
                     ),
+                    test_enum=hl.tstr,
                 ),
                 globals=hl.Struct(version='a_version'),
+                key=['locus', 'alleles'],
             ),
             hl.Table.parallelize(
                 [],
@@ -68,8 +70,10 @@ def test_create_globals_from_dataset_ht_configs(
                         test_select=hl.tint,
                         test_select_2=hl.tint,
                     ),
+                    test_enum=hl.tstr,
                 ),
                 globals=hl.Struct(version='b_version'),
+                key=['locus', 'alleles'],
             ),
         ]
         dataset_config_globals = Globals.from_dataset_configs(
@@ -89,7 +93,7 @@ def test_create_globals_from_dataset_ht_configs(
         )
         self.assertTrue(
             dataset_config_globals.selects
-            == {'a': {'test_select'}, 'b': {'test_select', 'field2'}},
+            == {'a': {'test_select', 'test_enum_id'}, 'b': {'test_select', 'field2', 'test_enum_id'}},
         )
 
     def test_from_rdc_or_annotations_ht(self):
@@ -115,7 +119,7 @@ def test_from_rdc_or_annotations_ht(self):
                     screen='v2',
                 ),
                 enums=hl.Struct(
-                    screen={'region_type': ['C', 'D']},
+                    screen=hl.Struct(region_type=['C', 'D']),
                 ),
             ),
         )

v03_pipeline/lib/reference_data/dataset_table_operations.py

@@ -20,11 +20,21 @@ def update_or_create_joined_ht(
     joined_ht: hl.Table,
 ) -> hl.Table:
     for dataset in datasets:
-        dataset_ht = get_dataset_ht(dataset, reference_genome)
-
+        # Drop the dataset if it exists.
         if dataset in joined_ht.row:
             joined_ht = joined_ht.drop(dataset)
+            joined_ht = joined_ht.annotate_globals(
+                paths=joined_ht.paths.drop(dataset),
+                versions=joined_ht.versions.drop(dataset),
+                enums=joined_ht.enums.drop(dataset),
+            )
 
+        # Handle cases where a dataset has been dropped OR renamed.
+        if dataset not in CONFIG:
+            continue
+
+        # Join the new one!
+        dataset_ht = get_dataset_ht(dataset, reference_genome)
         joined_ht = joined_ht.join(dataset_ht, 'outer')
         joined_ht = annotate_dataset_globals(joined_ht, dataset, dataset_ht)
 
@@ -51,7 +61,7 @@ def get_dataset_ht(
 
     ht = ht.filter(config['filter'](ht)) if 'filter' in config else ht
     ht = ht.select(**get_all_select_fields(ht, config))
-    ht = ht.transmute(**get_enum_select_fields(config.get('enum_select'), ht))
+    ht = ht.transmute(**get_enum_select_fields(ht, config))
     ht = ht.select_globals(
         path=(config['source_path'] if 'custom_import' in config else config['path']),
         version=parse_dataset_version(ht, dataset, config),
@@ -129,7 +139,8 @@ def get_all_select_fields(
     }
 
 
-def get_enum_select_fields(enum_selects: dict | None, ht: hl.Table) -> dict:
+def get_enum_select_fields(ht: hl.Table, config: dict) -> dict:
+    enum_selects = config.get('enum_select')
     enum_select_fields = {}
     if enum_selects is None:
         return enum_select_fields

v03_pipeline/lib/reference_data/dataset_table_operations_test.py

@@ -244,11 +244,13 @@ def test_get_enum_select_fields(self):
             ),
         )
         enum_select_fields = get_enum_select_fields(
+            ht,
             {
-                'variant': ['1', '2', '3', '4'],
-                'sv_type': ['a', 'b', 'c', 'd'],
+                'enum_select': {
+                    'variant': ['1', '2', '3', '4'],
+                    'sv_type': ['a', 'b', 'c', 'd'],
+                },
             },
-            ht,
         )
         mapped_ht = ht.transmute(**enum_select_fields)
         self.assertListEqual(
@@ -262,8 +264,10 @@ def test_get_enum_select_fields(self):
         )
 
         enum_select_fields = get_enum_select_fields(
-            {'sv_type': ['d']},
             ht,
+            {
+                'enum_select': {'sv_type': ['d']},
+            },
         )
         mapped_ht = ht.select(**enum_select_fields)
         self.assertRaises(Exception, mapped_ht.collect)
@@ -472,6 +476,7 @@ def test_parse_version(self, mock_read_table):
         )
         self.assertRaises(Exception, ht.globals.collect)
 
+    @mock.patch.dict(f'{PATH_TO_FILE_UNDER_TEST}.CONFIG', MOCK_CONFIG)
     @mock.patch(f'{PATH_TO_FILE_UNDER_TEST}.get_dataset_ht')
     @mock.patch(f'{PATH_TO_FILE_UNDER_TEST}.datetime', wraps=datetime)
     @mock.patch.object(ReferenceDatasetCollection, 'datasets')
@@ -538,3 +543,44 @@ def test_update_or_create_joined_ht_all_datasets(
             EXPECTED_JOINED_DATA,
         )
         self.assertCountEqual(ht.globals.collect(), EXPECTED_GLOBALS)
+
+    @mock.patch.dict(f'{PATH_TO_FILE_UNDER_TEST}.CONFIG', MOCK_CONFIG)
+    @mock.patch.object(ReferenceDatasetCollection, 'datasets')
+    def test_update_or_create_joined_ht_drop_a_dataset(
+        self,
+        mock_reference_dataset_collection_datasets,
+    ):
+        mock_reference_dataset_collection_datasets.return_value = ['b']
+        ht = hl.Table.parallelize(
+            [],
+            hl.tstruct(
+                locus=hl.tlocus(ReferenceGenome.GRCh38.value),
+                alleles=hl.tarray(hl.tstr),
+                c=hl.tint32,
+                b=hl.tint32,
+            ),
+            key=('locus', 'alleles'),
+            globals=hl.Struct(
+                paths=hl.Struct(c='abc', b='123'),
+                versions=hl.Struct(c='def', b='456'),
+                enums=hl.Struct(c=hl.Struct(d=['a', 'b'])),
+            ),
+        )
+        ht = update_or_create_joined_ht(
+            ReferenceDatasetCollection.COMBINED,
+            DatasetType.SNV_INDEL,
+            ReferenceGenome.GRCh38,
+            datasets=['c'],
+            joined_ht=ht,
+        )
+        self.assertCountEqual(
+            ht.globals.collect(),
+            [
+                hl.Struct(
+                    paths=hl.Struct(b='123'),
+                    versions=hl.Struct(b='456'),
+                    enums=hl.Struct(),
+                ),
+            ],
+        )
+

v03_pipeline/lib/tasks/base/base_variant_annotations_table.py

@@ -1,6 +1,7 @@
 import hail as hl
 import luigi
 
+from v03_pipeline.lib.annotations.enums import annotate_enums
 from v03_pipeline.lib.model import Env, ReferenceDatasetCollection
 from v03_pipeline.lib.paths import (
     valid_reference_dataset_collection_path,
@@ -79,10 +80,15 @@ def initialize_table(self) -> hl.Table:
     def update_table(self, ht: hl.Table) -> hl.Table:
         return ht
 
-    def annotate_reference_dataset_collection_globals(
+    def annotate_globals(
         self,
         ht: hl.Table,
     ) -> hl.Table:
+        ht = ht.annotate_globals(
+            paths=hl.Struct(),
+            versions=hl.Struct(),
+            enums=hl.Struct(),
+        )
         for rdc in ReferenceDatasetCollection.for_reference_genome_dataset_type(
             self.reference_genome,
             self.dataset_type,
@@ -104,4 +110,4 @@ def annotate_reference_dataset_collection_globals(
                 ),
                 updates=ht.globals.updates,
             )
-        return ht
+        return annotate_enums(ht, self.reference_genome, self.dataset_type)

v03_pipeline/lib/tasks/reference_data/update_variant_annotations_table_with_updated_reference_dataset.py

@@ -1,5 +1,6 @@
 import hail as hl
 
+from v03_pipeline.lib.annotations.fields import get_fields
 from v03_pipeline.lib.model import ReferenceDatasetCollection
 from v03_pipeline.lib.reference_data.compare_globals import (
     Globals,
@@ -13,18 +14,17 @@
 class UpdateVariantAnnotationsTableWithUpdatedReferenceDataset(
     BaseVariantAnnotationsTableTask,
 ):
-    _datasets_to_update: list[str]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._datasets_to_update = []
 
     @property
     def reference_dataset_collections(self) -> list[ReferenceDatasetCollection]:
-        return [
-            rdc
-            for rdc in ReferenceDatasetCollection.for_reference_genome_dataset_type(
-                self.reference_genome,
-                self.dataset_type,
-            )
-            if not rdc.requires_annotation
-        ]
+        return ReferenceDatasetCollection.for_reference_genome_dataset_type(
+            self.reference_genome,
+            self.dataset_type,
+        )
 
     def complete(self) -> bool:
         self._datasets_to_update = []
@@ -63,10 +63,18 @@ def update_table(self, ht: hl.Table) -> hl.Table:
         for dataset in self._datasets_to_update:
             rdc = ReferenceDatasetCollection.for_dataset(dataset, self.dataset_type)
             rdc_ht = self.rdc_annotation_dependencies[f'{rdc.value}_ht']
-
             if dataset in ht.row:
                 ht = ht.drop(dataset)
-
-            ht = ht.join(rdc_ht.select(dataset), 'left')
-
-        return self.annotate_reference_dataset_collection_globals(ht)
+            if rdc.requires_annotation:
+                formatting_fn = next(x for x in self.dataset_type.formatting_annotation_fns(self.reference_genome) if x.__name__ == dataset)
+                ht = ht.annotate(
+                    **get_fields(
+                        ht,
+                        [formatting_fn],
+                        **self.rdc_annotation_dependencies,
+                        **self.param_kwargs,
+                    ),
+                )
+            else:
+                ht = ht.join(rdc_ht.select(dataset), 'left')
+        return self.annotate_globals(ht)