new task for updating annotations w/ rdc data (#676)

* refactor globals and create new task to update vat

* see if tests pass:

* test for new task

* getting Unfulfilled dependency at runtime

* load_gencode

* fix 1 test

* a few more tests for the task

* big test

* tests

* oops

* most tests should pass.

* do not clear globals if updating annotations table with rdc update

* fix da tests

* pr comments

* oops

* refactor

* honestly cleaner

* even cleaner

Author: jklugherz

v03_pipeline/lib/model/reference_dataset_collection.py

@@ -94,3 +94,16 @@ def for_reference_genome_dataset_type(
         if not Env.ACCESS_PRIVATE_REFERENCE_DATASETS:
             return [rdc for rdc in rdcs if rdc.access_control == AccessControl.PUBLIC]
         return rdcs
+
+    @classmethod
+    def for_dataset(
+        cls,
+        dataset: str,
+        dataset_type: DatasetType,
+    ) -> 'ReferenceDatasetCollection':
+        for rdc in cls:
+            if dataset in rdc.datasets(dataset_type):
+                return rdc
+
+        err_msg = f'Dataset "{dataset}" not found in any reference dataset collection'
+        raise ValueError(err_msg)

v03_pipeline/lib/tasks/base/base_variant_annotations_table.py

@@ -7,10 +7,29 @@
     variant_annotations_table_path,
 )
 from v03_pipeline.lib.tasks.base.base_update_task import BaseUpdateTask
-from v03_pipeline.lib.tasks.files import GCSorLocalTarget, HailTableTask
+from v03_pipeline.lib.tasks.files import GCSorLocalTarget
+from v03_pipeline.lib.tasks.reference_data.updated_reference_dataset_collection import (
+    UpdatedReferenceDatasetCollectionTask,
+)
 
 
 class BaseVariantAnnotationsTableTask(BaseUpdateTask):
+    @property
+    def rdc_annotation_dependencies(self) -> dict[str, hl.Table]:
+        annotation_dependencies = {}
+        for rdc in ReferenceDatasetCollection.for_reference_genome_dataset_type(
+            self.reference_genome,
+            self.dataset_type,
+        ):
+            annotation_dependencies[f'{rdc.value}_ht'] = hl.read_table(
+                valid_reference_dataset_collection_path(
+                    self.reference_genome,
+                    self.dataset_type,
+                    rdc,
+                ),
+            )
+        return annotation_dependencies
+
     def output(self) -> luigi.Target:
         return GCSorLocalTarget(
             variant_annotations_table_path(
@@ -21,12 +40,11 @@ def output(self) -> luigi.Target:
 
     def requires(self) -> list[luigi.Task]:
         return [
-            HailTableTask(
-                valid_reference_dataset_collection_path(
-                    self.reference_genome,
-                    self.dataset_type,
-                    rdc,
-                ),
+            UpdatedReferenceDatasetCollectionTask(
+                self.reference_genome,
+                self.dataset_type,
+                self.sample_type,
+                rdc,
             )
             for rdc in ReferenceDatasetCollection.for_reference_genome_dataset_type(
                 self.reference_genome,
@@ -50,3 +68,30 @@ def initialize_table(self) -> hl.Table:
 
     def update_table(self, ht: hl.Table) -> hl.Table:
         return ht
+
+    def annotate_reference_dataset_collection_globals(
+        self,
+        ht: hl.Table,
+    ) -> hl.Table:
+        for rdc in ReferenceDatasetCollection.for_reference_genome_dataset_type(
+            self.reference_genome,
+            self.dataset_type,
+        ):
+            rdc_ht = self.rdc_annotation_dependencies[f'{rdc.value}_ht']
+            rdc_globals = rdc_ht.index_globals()
+            ht = ht.select_globals(
+                paths=hl.Struct(
+                    **ht.globals.paths,
+                    **rdc_globals.paths,
+                ),
+                versions=hl.Struct(
+                    **ht.globals.versions,
+                    **rdc_globals.versions,
+                ),
+                enums=hl.Struct(
+                    **ht.globals.enums,
+                    **rdc_globals.enums,
+                ),
+                updates=ht.globals.updates,
+            )
+        return ht

v03_pipeline/lib/tasks/base/base_variant_annotations_table_test.py

@@ -1,4 +1,5 @@
 import shutil
+from unittest.mock import patch
 
 import hail as hl
 import luigi.worker
@@ -14,6 +15,7 @@
     BaseVariantAnnotationsTableTask,
 )
 from v03_pipeline.lib.tasks.files import GCSorLocalFolderTarget
+from v03_pipeline.lib.test.mock_complete_task import MockCompleteTask
 from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
 
 TEST_COMBINED_1 = 'v03_pipeline/var/test/reference_data/test_combined_1.ht'
@@ -49,7 +51,11 @@ def setUp(self) -> None:
             ),
         )
 
-    def test_should_create_initialized_table(self) -> None:
+    @patch(
+        'v03_pipeline.lib.tasks.base.base_variant_annotations_table.UpdatedReferenceDatasetCollectionTask',
+    )
+    def test_should_create_initialized_table(self, mock_update_rdc_task) -> None:
+        mock_update_rdc_task.return_value = MockCompleteTask()
         vat_task = BaseVariantAnnotationsTableTask(
             reference_genome=ReferenceGenome.GRCh38,
             dataset_type=DatasetType.SNV_INDEL,

v03_pipeline/lib/tasks/reference_data/update_variant_annotations_table_with_updated_reference_dataset.py

@@ -0,0 +1,72 @@
+import hail as hl
+
+from v03_pipeline.lib.model import ReferenceDatasetCollection
+from v03_pipeline.lib.reference_data.compare_globals import (
+    Globals,
+    get_datasets_to_update,
+)
+from v03_pipeline.lib.tasks.base.base_variant_annotations_table import (
+    BaseVariantAnnotationsTableTask,
+)
+
+
+class UpdateVariantAnnotationsTableWithUpdatedReferenceDataset(
+    BaseVariantAnnotationsTableTask,
+):
+    _datasets_to_update: list[str]
+
+    @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
+        ]
+
+    def complete(self) -> bool:
+        self._datasets_to_update = []
+
+        if not super().complete():
+            for rdc in self.reference_dataset_collections:
+                self._datasets_to_update.extend(
+                    rdc.datasets(
+                        self.dataset_type,
+                    ),
+                )
+            return False
+
+        for rdc in self.reference_dataset_collections:
+            annotations_ht_globals = Globals.from_ht(
+                hl.read_table(self.output().path),
+                rdc,
+                self.dataset_type,
+            )
+            rdc_ht_globals = Globals.from_ht(
+                self.rdc_annotation_dependencies[f'{rdc.value}_ht'],
+                rdc,
+                self.dataset_type,
+            )
+            self._datasets_to_update.extend(
+                get_datasets_to_update(
+                    rdc,
+                    annotations_ht_globals,
+                    rdc_ht_globals,
+                    self.dataset_type,
+                ),
+            )
+        return not self._datasets_to_update
+
+    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)