Merge branch 'dev' of github.com:broadinstitute/seqr-loading-pipelines

Author: bpblanken

v03_pipeline/lib/misc/family_entries.py

@@ -107,7 +107,7 @@ def remove_new_callset_family_guids(
         family_entries=(
             hl.array(family_indexes_to_keep).map(lambda i: ht.family_entries[i])
             if len(family_indexes_to_keep) > 0
-            else hl.missing(ht.family_entries.dtype.element_type)
+            else hl.missing(ht.family_entries.dtype)
         ),
     )
     return ht.annotate_globals(

v03_pipeline/lib/reference_data/compare_globals.py

@@ -4,8 +4,6 @@
 
 from v03_pipeline.lib.logger import get_logger
 from v03_pipeline.lib.model import (
-    DatasetType,
-    ReferenceDatasetCollection,
     ReferenceGenome,
 )
 from v03_pipeline.lib.reference_data.config import CONFIG
@@ -22,8 +20,8 @@
 
 @dataclasses.dataclass
 class Globals:
-    paths: dict[str]
-    versions: dict[str]
+    paths: dict[str, str]
+    versions: dict[str, str]
     enums: dict[str, dict[str, list[str]]]
     selects: dict[str, set[str]]
 
@@ -33,12 +31,11 @@ def __getitem__(self, name: str):
     @classmethod
     def from_dataset_configs(
         cls,
-        rdc: ReferenceDatasetCollection,
-        dataset_type: DatasetType,
         reference_genome: ReferenceGenome,
+        datasets: list[str],
     ):
         paths, versions, enums, selects = {}, {}, {}, {}
-        for dataset in rdc.datasets(dataset_type):
+        for dataset in datasets:
             dataset_config = CONFIG[dataset][reference_genome.v02_value]
             dataset_ht = import_ht_from_config_path(dataset_config, reference_genome)
 
@@ -64,17 +61,21 @@ def from_dataset_configs(
     def from_ht(
         cls,
         ht: hl.Table,
-        rdc: ReferenceDatasetCollection,
-        dataset_type: DatasetType,
+        datasets: list[str],
     ):
         rdc_globals_struct = hl.eval(ht.globals)
         paths = dict(rdc_globals_struct.paths)
         versions = dict(rdc_globals_struct.versions)
         # enums are nested structs
         enums = {k: dict(v) for k, v in rdc_globals_struct.enums.items()}
 
+        for global_dict in [paths, versions, enums]:
+            for dataset in list(global_dict.keys()):
+                if dataset not in datasets:
+                    global_dict.pop(dataset)
+
         selects = {}
-        for dataset in rdc.datasets(dataset_type):
+        for dataset in datasets:
             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] = (
@@ -86,30 +87,20 @@ def from_ht(
 
 
 def get_datasets_to_update(
-    rdc: ReferenceDatasetCollection,
     ht1_globals: Globals,
     ht2_globals: Globals,
-    dataset_type: DatasetType,
 ) -> list[str]:
-    return [
-        dataset
-        for dataset in rdc.datasets(dataset_type)
-        if not validate_globals_match(rdc, ht1_globals, ht2_globals, dataset)
-    ]
-
+    datasets_to_update = set()
 
-def validate_globals_match(
-    rdc: ReferenceDatasetCollection,
-    ht1_globals: Globals,
-    ht2_globals: Globals,
-    dataset: str,
-) -> bool:
-    results = []
     for field in dataclasses.fields(Globals):
-        result = ht1_globals[field.name].get(dataset) == ht2_globals[field.name].get(
-            dataset,
+        datasets_to_update.update(
+            ht1_globals[field.name].keys() ^ ht2_globals[field.name].keys(),
         )
-        if result is False:
-            logger.info(f'{field.name} mismatch for {dataset}, {rdc.value}')
-        results.append(result)
-    return all(results)
+        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,
+            ):
+                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

@@ -4,55 +4,68 @@
 import hail as hl
 
 from v03_pipeline.lib.model import (
-    DatasetType,
-    ReferenceDatasetCollection,
     ReferenceGenome,
 )
 from v03_pipeline.lib.reference_data.compare_globals import (
     Globals,
     get_datasets_to_update,
 )
 
-
-class CompareGlobalsTest(unittest.TestCase):
-    @mock.patch.dict(
-        'v03_pipeline.lib.reference_data.compare_globals.CONFIG',
-        {
-            'a': {
-                '38': {
-                    'custom_import': None,
-                    'source_path': 'a_path',  # 'a' has a custom import
-                    'select': {
-                        'test_select': 'info.test_select',
-                        'test_enum': 'test_enum',
-                    },
-                    'version': 'a_version',
-                    'enum_select': {'test_enum': ['A', 'B']},
-                },
+CONFIG = {
+    'a': {
+        '38': {
+            'custom_import': None,
+            'source_path': 'a_path',  # 'a' has a custom import
+            'select': {
+                'test_select': 'info.test_select',
+                'test_enum': 'test_enum',
             },
-            'b': {  # b is missing version
-                '38': {
-                    'path': 'b_path',
-                    '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},
-                },
+            'version': 'a_version',
+            'enum_select': {'test_enum': ['A', 'B']},
+        },
+    },
+    'b': {  # b is missing version
+        '38': {
+            'path': 'b_path',
+            '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},
         },
-    )
+    },
+}
+
+B_TABLE = hl.Table.parallelize(
+    [],
+    schema=hl.tstruct(
+        locus=hl.tlocus('GRCh38'),
+        alleles=hl.tarray(hl.tstr),
+        info=hl.tstruct(
+            test_select=hl.tint,
+            test_select_2=hl.tint,
+        ),
+        test_enum=hl.tstr,
+    ),
+    globals=hl.Struct(
+        version='b_version',
+        path='b_path',
+        enums=hl.Struct(test_enum=['C', 'D']),
+    ),
+    key=['locus', 'alleles'],
+)
+
+
+class CompareGlobalsTest(unittest.TestCase):
+    @mock.patch.dict('v03_pipeline.lib.reference_data.compare_globals.CONFIG', CONFIG)
     @mock.patch(
         'v03_pipeline.lib.reference_data.compare_globals.import_ht_from_config_path',
     )
-    @mock.patch.object(ReferenceDatasetCollection, 'datasets')
-    def test_create_globals_from_dataset_ht_configs(
+    def test_create_globals_from_dataset_configs(
         self,
-        mock_rdc_datasets,
         mock_import_dataset_ht,
     ):
-        mock_rdc_datasets.return_value = ['a', 'b']
         mock_import_dataset_ht.side_effect = [
             hl.Table.parallelize(
                 [],
@@ -64,28 +77,18 @@ def test_create_globals_from_dataset_ht_configs(
                     ),
                     test_enum=hl.tstr,
                 ),
-                globals=hl.Struct(version='a_version'),
-                key=['locus', 'alleles'],
-            ),
-            hl.Table.parallelize(
-                [],
-                schema=hl.tstruct(
-                    locus=hl.tlocus('GRCh38'),
-                    alleles=hl.tarray(hl.tstr),
-                    info=hl.tstruct(
-                        test_select=hl.tint,
-                        test_select_2=hl.tint,
-                    ),
-                    test_enum=hl.tstr,
+                globals=hl.Struct(
+                    version='a_version',
+                    path='a_path',
+                    enums=hl.Struct(test_enum=['A', 'B']),
                 ),
-                globals=hl.Struct(version='b_version'),
                 key=['locus', 'alleles'],
             ),
+            B_TABLE,
         ]
         dataset_config_globals = Globals.from_dataset_configs(
-            rdc=ReferenceDatasetCollection.INTERVAL,
-            dataset_type=DatasetType.SNV_INDEL,
             reference_genome=ReferenceGenome.GRCh38,
+            datasets=['a', 'b'],
         )
         self.assertTrue(
             dataset_config_globals.versions == {'a': 'a_version', 'b': 'b_version'},
@@ -105,6 +108,36 @@ def test_create_globals_from_dataset_ht_configs(
             },
         )
 
+    @mock.patch.dict('v03_pipeline.lib.reference_data.compare_globals.CONFIG', CONFIG)
+    @mock.patch(
+        'v03_pipeline.lib.reference_data.dataset_table_operations.hl.read_table',
+    )
+    def test_create_globals_from_dataset_configs_single_dataset(self, mock_read_table):
+        # by mocking hl.read_table() (only possible for a dataset without a custom import),
+        # we can test the code inside import_ht_from_config_path()
+        mock_read_table.return_value = B_TABLE
+
+        dataset_config_globals = Globals.from_dataset_configs(
+            reference_genome=ReferenceGenome.GRCh38,
+            datasets=['b'],
+        )
+
+        self.assertTrue(
+            dataset_config_globals.versions == {'b': 'b_version'},
+        )
+        self.assertTrue(
+            dataset_config_globals.paths == {'b': 'b_path'},
+        )
+        self.assertTrue(
+            dataset_config_globals.enums == {'b': {'test_enum': ['C', 'D']}},
+        )
+        self.assertTrue(
+            dataset_config_globals.selects
+            == {
+                'b': {'test_select', 'field2', 'test_enum_id'},
+            },
+        )
+
     def test_from_rdc_or_annotations_ht(self):
         rdc_ht = hl.Table.parallelize(
             [],
@@ -134,8 +167,7 @@ def test_from_rdc_or_annotations_ht(self):
         )
         rdc_globals = Globals.from_ht(
             rdc_ht,
-            rdc=ReferenceDatasetCollection.INTERVAL,
-            dataset_type=DatasetType.SNV_INDEL,
+            ['gnomad_non_coding_constraint', 'screen'],
         )
         self.assertTrue(
             rdc_globals.versions
@@ -159,11 +191,8 @@ def test_from_rdc_or_annotations_ht(self):
             },
         )
 
-    @mock.patch.object(ReferenceDatasetCollection, 'datasets')
-    def test_get_datasets_to_update_version_different(self, mock_rdc_datasets):
-        mock_rdc_datasets.return_value = ['a', 'b', 'c']
+    def test_get_datasets_to_update_version_different(self):
         result = get_datasets_to_update(
-            rdc=ReferenceDatasetCollection.INTERVAL,
             ht1_globals=Globals(
                 paths={'a': 'a_path', 'b': 'b_path'},
                 # 'a' has a different version, 'c' is missing version in ht2_globals
@@ -177,15 +206,11 @@ def test_get_datasets_to_update_version_different(self, mock_rdc_datasets):
                 enums={'a': {}, 'b': {}},
                 selects={'a': set(), 'b': set()},
             ),
-            dataset_type=DatasetType.SNV_INDEL,
         )
         self.assertTrue(result == ['a', 'c'])
 
-    @mock.patch.object(ReferenceDatasetCollection, 'datasets')
-    def test_get_datasets_to_update_path_different(self, mock_rdc_datasets):
-        mock_rdc_datasets.return_value = ['a', 'b', 'c']
+    def test_get_datasets_to_update_path_different(self):
         result = get_datasets_to_update(
-            rdc=ReferenceDatasetCollection.INTERVAL,
             ht1_globals=Globals(
                 # 'b' has a different path, 'c' is missing path in ht2_globals
                 paths={'a': 'a_path', 'b': 'old_b_path', 'c': 'extra_c_path'},
@@ -199,15 +224,11 @@ def test_get_datasets_to_update_path_different(self, mock_rdc_datasets):
                 enums={'a': {}, 'b': {}},
                 selects={'a': set(), 'b': set()},
             ),
-            dataset_type=DatasetType.SNV_INDEL,
         )
         self.assertTrue(result == ['b', 'c'])
 
-    @mock.patch.object(ReferenceDatasetCollection, 'datasets')
-    def test_get_datasets_to_update_enum_different(self, mock_rdc_datasets):
-        mock_rdc_datasets.return_value = ['a', 'b', 'c']
+    def test_get_datasets_to_update_enum_different(self):
         result = get_datasets_to_update(
-            rdc=ReferenceDatasetCollection.INTERVAL,
             ht1_globals=Globals(
                 paths={'a': 'a_path', 'b': 'b_path'},
                 versions={'a': 'v1', 'b': 'v2'},
@@ -225,15 +246,11 @@ def test_get_datasets_to_update_enum_different(self, mock_rdc_datasets):
                 enums={'a': {'test_enum': ['C', 'D']}, 'b': {'enum_key_2': []}},
                 selects={'a': set(), 'b': set()},
             ),
-            dataset_type=DatasetType.SNV_INDEL,
         )
         self.assertTrue(result == ['a', 'b', 'c'])
 
-    @mock.patch.object(ReferenceDatasetCollection, 'datasets')
-    def test_get_datasets_to_update_select_different(self, mock_rdc_datasets):
-        mock_rdc_datasets.return_value = ['a', 'b', 'c']
+    def test_get_datasets_to_update_select_different(self):
         result = get_datasets_to_update(
-            rdc=ReferenceDatasetCollection.INTERVAL,
             ht1_globals=Globals(
                 paths={'a': 'a_path', 'b': 'b_path'},
                 versions={'a': 'v1', 'b': 'v2'},
@@ -251,6 +268,5 @@ def test_get_datasets_to_update_select_different(self, mock_rdc_datasets):
                 enums={'a': {}, 'b': {}},
                 selects={'a': {'field1'}, 'b': {'test_select_2'}},
             ),
-            dataset_type=DatasetType.SNV_INDEL,
         )
         self.assertTrue(result == ['a', 'b', 'c'])