Reference Data Update Type Equality Check (#789)

* Finish validity check test

* ruff

* update dbnsfp field

* More types

* more types

* ugh

* twiddle it back

* update type

* more tweaks

* lint

* fix floats

* decompose

* ruff formatg

* Update compare_globals_test.py

Author: bpblanken

v03_pipeline/lib/reference_data/compare_globals.py

@@ -21,7 +21,7 @@ class Globals:
     paths: dict[str, str]
     versions: dict[str, str]
     enums: dict[str, dict[str, list[str]]]
-    selects: dict[str, set[str]]
+    selects: dict[str, dict[str, hl.dtype]]
 
     def __getitem__(self, name: str):
         return getattr(self, name)
@@ -50,7 +50,11 @@ def from_dataset_configs(
             dataset_ht = dataset_ht.transmute(
                 **get_enum_select_fields(dataset_ht, dataset_config),
             )
-            selects[dataset] = set(dataset_ht.row) - set(dataset_ht.key)
+            selects[dataset] = {
+                k: v.dtype
+                for k, v in dict(dataset_ht.row).items()
+                if k not in set(dataset_ht.key)
+            }
         return cls(paths, versions, enums, selects)
 
     @classmethod
@@ -69,32 +73,52 @@ def from_ht(
             if dataset in ht.row:
                 # NB: handle an edge case (mito high constraint) where we annotate a bool from the reference dataset collection
                 selects[dataset] = (
-                    set(ht[dataset])
+                    {k: v.dtype for k, v in dict(ht[dataset]).items()}
                     if isinstance(ht[dataset], hl.StructExpression)
-                    else set()
+                    else {}
                 )
         return cls(paths, versions, enums, selects)
 
 
+def validate_selects_types(
+    ht1_globals: Globals,
+    ht2_globals: Globals,
+    dataset: str,
+) -> None:
+    # Assert that all shared annotations have identical types
+    shared_selects = (
+        ht1_globals['selects'][dataset].keys()
+        & ht2_globals['selects'].get(dataset).keys()
+    )
+    mismatched_select_types = [
+        (select, ht2_globals['selects'][dataset][select])
+        for select in shared_selects
+        if (
+            ht1_globals['selects'][dataset][select]
+            != ht2_globals['selects'][dataset][select]
+        )
+    ]
+    if mismatched_select_types:
+        msg = f'Unexpected field types detected in {dataset}: {mismatched_select_types}'
+        raise ValueError(msg)
+
+
 def get_datasets_to_update(
     ht1_globals: Globals,
     ht2_globals: Globals,
     validate_selects: bool = True,
 ) -> list[str]:
     datasets_to_update = set()
-
     for field in dataclasses.fields(Globals):
         if field.name == 'selects' and not validate_selects:
             continue
-
         datasets_to_update.update(
             ht1_globals[field.name].keys() ^ ht2_globals[field.name].keys(),
         )
         for dataset in ht1_globals[field.name].keys() & ht2_globals[field.name].keys():
-            if ht1_globals[field.name].get(dataset) != ht2_globals[field.name].get(
-                dataset,
-            ):
+            if field.name == 'selects':
+                validate_selects_types(ht1_globals, ht2_globals, dataset)
+            if ht1_globals[field.name][dataset] != ht2_globals[field.name][dataset]:
                 logger.info(f'{field.name} mismatch for {dataset}')
                 datasets_to_update.add(dataset)
-
     return sorted(datasets_to_update)

v03_pipeline/lib/reference_data/compare_globals_test.py

@@ -103,8 +103,15 @@ def test_create_globals_from_dataset_configs(
         self.assertTrue(
             dataset_config_globals.selects
             == {
-                'a': {'test_select', 'test_enum_id'},
-                'b': {'test_select', 'field2', 'test_enum_id'},
+                'a': {
+                    'test_select': hl.tint32,
+                    'test_enum_id': hl.tint32,
+                },
+                'b': {
+                    'test_select': hl.tint32,
+                    'field2': hl.tint32,
+                    'test_enum_id': hl.tint32,
+                },
             },
         )
 
@@ -134,7 +141,11 @@ def test_create_globals_from_dataset_configs_single_dataset(self, mock_read_tabl
         self.assertTrue(
             dataset_config_globals.selects
             == {
-                'b': {'test_select', 'field2', 'test_enum_id'},
+                'b': {
+                    'test_select': hl.tint32,
+                    'field2': hl.tint32,
+                    'test_enum_id': hl.tint32,
+                },
             },
         )
 
@@ -186,8 +197,8 @@ def test_from_rdc_or_annotations_ht(self):
         self.assertTrue(
             rdc_globals.selects
             == {
-                'gnomad_non_coding_constraint': {'z_score'},
-                'screen': {'region_type_ids'},
+                'gnomad_non_coding_constraint': {'z_score': hl.tfloat32},
+                'screen': {'region_type_ids': hl.tarray(hl.tint32)},
             },
         )
 
@@ -198,13 +209,13 @@ def test_get_datasets_to_update_version_different(self):
                 # 'a' has a different version, 'c' is missing version in ht2_globals
                 versions={'a': 'v2', 'b': 'v2', 'c': 'v1'},
                 enums={'a': {}, 'b': {}, 'c': {}},
-                selects={'a': set(), 'b': set()},
+                selects={'a': {}, 'b': {}},
             ),
             ht2_globals=Globals(
                 paths={'a': 'a_path', 'b': 'b_path'},
                 versions={'a': 'v1', 'b': 'v2'},
                 enums={'a': {}, 'b': {}},
-                selects={'a': set(), 'b': set()},
+                selects={'a': {}, 'b': {}},
             ),
         )
         self.assertTrue(result == ['a', 'c'])
@@ -216,13 +227,13 @@ def test_get_datasets_to_update_path_different(self):
                 paths={'a': 'a_path', 'b': 'old_b_path', 'c': 'extra_c_path'},
                 versions={'a': 'v1', 'b': 'v2'},
                 enums={'a': {}, 'b': {}},
-                selects={'a': set(), 'b': set()},
+                selects={'a': {}, 'b': {}},
             ),
             ht2_globals=Globals(
                 paths={'a': 'a_path', 'b': 'b_path'},
                 versions={'a': 'v1', 'b': 'v2'},
                 enums={'a': {}, 'b': {}},
-                selects={'a': set(), 'b': set()},
+                selects={'a': {}, 'b': {}},
             ),
         )
         self.assertTrue(result == ['b', 'c'])
@@ -238,13 +249,13 @@ def test_get_datasets_to_update_enum_different(self):
                     'b': {'enum_key_1': []},
                     'c': {},
                 },
-                selects={'a': set(), 'b': set()},
+                selects={'a': {}, 'b': {}},
             ),
             ht2_globals=Globals(
                 paths={'a': 'a_path', 'b': 'b_path'},
                 versions={'a': 'v1', 'b': 'v2'},
                 enums={'a': {'test_enum': ['C', 'D']}, 'b': {'enum_key_2': []}},
-                selects={'a': set(), 'b': set()},
+                selects={'a': {}, 'b': {}},
             ),
         )
         self.assertTrue(result == ['a', 'b', 'c'])
@@ -257,16 +268,54 @@ def test_get_datasets_to_update_select_different(self):
                 enums={'a': {}, 'b': {}},
                 # 'a' has extra select, 'b' has different select, 'c' is missing select in ht2_globals
                 selects={
-                    'a': {'field1', 'field2'},
-                    'b': {'test_select'},
-                    'c': set('test_select'),
+                    'a': {'field1': hl.tint32, 'field2': hl.tint32},
+                    'b': {'test_select': hl.tint32},
+                    'c': {'test_select': hl.tint32},
                 },
             ),
             ht2_globals=Globals(
                 paths={'a': 'a_path', 'b': 'b_path'},
                 versions={'a': 'v1', 'b': 'v2'},
                 enums={'a': {}, 'b': {}},
-                selects={'a': {'field1'}, 'b': {'test_select_2'}},
+                selects={'a': {'field1': hl.tint32}, 'b': {'test_select_2': hl.tint32}},
             ),
         )
         self.assertTrue(result == ['a', 'b', 'c'])
+
+    def test_get_datasets_to_update_select_type_validation(self):
+        self.assertRaisesRegex(
+            ValueError,
+            "Unexpected field types detected in a: \\[\\('field1', dtype\\('int32'\\)\\)\\]",
+            get_datasets_to_update,
+            ht1_globals=Globals(
+                paths={'a': 'a_path'},
+                versions={'a': 'v1'},
+                enums={'a': {}},
+                selects={
+                    'a': {'field1': hl.tarray(hl.tint32)},
+                },
+            ),
+            ht2_globals=Globals(
+                paths={'a': 'a_path'},
+                versions={'a': 'v1'},
+                enums={'a': {}},
+                selects={'a': {'field1': hl.tint32, 'field2': hl.tint32}},
+            ),
+        )
+        result = get_datasets_to_update(
+            ht1_globals=Globals(
+                paths={'a': 'a_path'},
+                versions={'a': 'v1'},
+                enums={'a': {}},
+                selects={
+                    'a': {'field1': hl.tarray(hl.tint32)},
+                },
+            ),
+            ht2_globals=Globals(
+                paths={'a': 'a_path'},
+                versions={'a': 'v1'},
+                enums={'a': {}},
+                selects={'a': {'field1': hl.tarray(hl.tint32), 'field2': hl.tint32}},
+            ),
+        )
+        self.assertTrue(result == ['a'])

v03_pipeline/lib/tasks/reference_data/update_variant_annotations_table_with_updated_reference_dataset_test.py

@@ -45,7 +45,7 @@
         hl.tstruct(
             locus=hl.tlocus('GRCh38'),
             alleles=hl.tarray(hl.tstr),
-            PHRED=hl.tint32,
+            PHRED=hl.tfloat32,
         ),
         key=['locus', 'alleles'],
         globals=hl.Struct(
@@ -760,7 +760,7 @@ def test_update_vat_with_updated_rdc_snv_indel_38(
                         conditions=None,
                     ),
                     dbnsfp=hl.Struct(
-                        REVEL_score=0.043,
+                        REVEL_score=0.0430000014603138,
                         SIFT_score=None,
                         Polyphen2_HVAR_score=None,
                         MutationTaster_pred_id=0,
@@ -1168,7 +1168,7 @@ def test_update_vat_with_updated_rdc_snv_indel_37(
                         conditions=None,
                     ),
                     dbnsfp=hl.Struct(
-                        REVEL_score=0.043,
+                        REVEL_score=0.0430000014603138,
                         SIFT_score=None,
                         Polyphen2_HVAR_score=None,
                         MutationTaster_pred_id=0,

v03_pipeline/var/test/reference_data/test_combined_1.ht/.README.txt.crc

@@ -1,3 +1,3 @@
 This folder comprises a Hail (www.hail.is) native Table or MatrixTable.
-  Written with version 0.2.128-eead8100a1c1
-  Created at 2024/05/09 20:02:21
+  Written with version 0.2.130-bea04d9c79b5
+  Created at 2024/05/20 13:48:16

v03_pipeline/var/test/reference_data/test_combined_1.ht/.metadata.json.gz.crc

@@ -1,3 +1,3 @@
 This folder comprises a Hail (www.hail.is) native Table or MatrixTable.
-  Written with version 0.2.128-eead8100a1c1
-  Created at 2024/03/21 11:28:13
+  Written with version 0.2.130-bea04d9c79b5
+  Created at 2024/05/20 15:38:26

v03_pipeline/var/test/reference_data/test_combined_1.ht/README.txt

@@ -1,3 +1,3 @@
 This folder comprises a Hail (www.hail.is) native Table or MatrixTable.
-  Written with version 0.2.128-eead8100a1c1
-  Created at 2024/03/21 11:35:30
+  Written with version 0.2.130-bea04d9c79b5
+  Created at 2024/05/20 14:08:17

v03_pipeline/var/test/reference_data/test_combined_1.ht/index/part-0-963ef6be-cd5f-443a-970c-8b5cf2bcd090.idx/.index.crc renamed to v03_pipeline/var/test/reference_data/test_combined_1.ht/index/part-0-6a5a9d6a-4ded-424b-9735-922a5346e7ad.idx/.index.crc

@@ -1,3 +1,3 @@
 This folder comprises a Hail (www.hail.is) native Table or MatrixTable.
-  Written with version 0.2.114-cc8d36408b36
-  Created at 2023/07/13 19:51:12
+  Written with version 0.2.130-bea04d9c79b5
+  Created at 2024/05/20 13:22:32