Merge pull request #762 from broadinstitute/dev

Dev

Author: bpblanken

v03_pipeline/lib/annotations/rdc_dependencies.py

@@ -0,0 +1,29 @@
+import hail as hl
+
+from v03_pipeline.lib.model import (
+    DatasetType,
+    ReferenceDatasetCollection,
+    ReferenceGenome,
+)
+from v03_pipeline.lib.paths import (
+    valid_reference_dataset_collection_path,
+)
+
+
+def get_rdc_annotation_dependencies(
+    dataset_type: DatasetType,
+    reference_genome: ReferenceGenome,
+) -> dict[str, hl.Table]:
+    deps = {}
+    for rdc in ReferenceDatasetCollection.for_reference_genome_dataset_type(
+        reference_genome,
+        dataset_type,
+    ):
+        deps[f'{rdc.value}_ht'] = hl.read_table(
+            valid_reference_dataset_collection_path(
+                reference_genome,
+                dataset_type,
+                rdc,
+            ),
+        )
+    return deps

v03_pipeline/lib/misc/callsets.py

@@ -0,0 +1,70 @@
+import functools
+
+import hail as hl
+
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome
+from v03_pipeline.lib.paths import remapped_and_subsetted_callset_path
+
+
+def get_callset_ht(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    callset_paths: list[str],
+    project_guids: list[str],
+    project_remap_paths: list[str],
+    project_pedigree_paths: list[str],
+):
+    callset_hts = [
+        hl.read_matrix_table(
+            remapped_and_subsetted_callset_path(
+                reference_genome,
+                dataset_type,
+                callset_path,
+                project_guid,
+            ),
+        ).rows()
+        for (callset_path, project_guid, _, _) in callset_project_pairs(
+            callset_paths,
+            project_guids,
+            project_remap_paths,
+            project_pedigree_paths,
+        )
+    ]
+
+    # Drop any fields potentially unshared/unused by the annotations.
+    for i, callset_ht in enumerate(callset_hts):
+        for row_field in dataset_type.optional_row_fields:
+            if hasattr(callset_ht, row_field):
+                callset_hts[i] = callset_ht.drop(row_field)
+
+    callset_ht = functools.reduce(
+        (lambda ht1, ht2: ht1.union(ht2, unify=True)),
+        callset_hts,
+    )
+    return callset_ht.distinct()
+
+
+def callset_project_pairs(
+    callset_paths: list[str],
+    project_guids: list[str],
+    project_remap_paths: list[str],
+    project_pedigree_paths: list[str],
+):
+    if len(callset_paths) == len(project_guids):
+        return zip(
+            callset_paths,
+            project_guids,
+            project_remap_paths,
+            project_pedigree_paths,
+            strict=True,
+        )
+    return (
+        (callset_path, project_guid, project_remap_path, project_pedigree_path)
+        for callset_path in callset_paths
+        for (project_guid, project_remap_path, project_pedigree_path) in zip(
+            project_guids,
+            project_remap_paths,
+            project_pedigree_paths,
+            strict=True,
+        )
+    )

v03_pipeline/lib/misc/util.py

@@ -236,3 +236,18 @@ def variant_annotations_table_path(
         ),
         'annotations.ht',
     )
+
+
+def new_variants_table_path(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    run_id: str,
+) -> str:
+    return os.path.join(
+        runs_path(
+            reference_genome,
+            dataset_type,
+        ),
+        run_id,
+        'new_variants.ht',
+    )

v03_pipeline/lib/paths.py

@@ -12,6 +12,7 @@
     imported_callset_path,
     lookup_table_path,
     metadata_for_run_path,
+    new_variants_table_path,
     project_table_path,
     relatedness_check_table_path,
     remapped_and_subsetted_callset_path,
@@ -160,3 +161,13 @@ def test_imported_callset_path(self) -> None:
             ),
             '/seqr-loading-temp/v03/GRCh38/SNV_INDEL/imported_callsets/ead56bb177a5de24178e1e622ce1d8beb3f8892bdae1c925d22ca0af4013d6dd.mt',
         )
+
+    def test_new_variants_table_path(self) -> None:
+        self.assertEqual(
+            new_variants_table_path(
+                ReferenceGenome.GRCh38,
+                DatasetType.SNV_INDEL,
+                'manual__2023-06-26T18:30:09.349671+00:00',
+            ),
+            '/hail-search-data/v03/GRCh38/SNV_INDEL/runs/manual__2023-06-26T18:30:09.349671+00:00/new_variants.ht',
+        )

v03_pipeline/lib/paths_test.py

@@ -1,4 +1,3 @@
-from v03_pipeline.lib.tasks.delete_old_runs import DeleteOldRunsTask
 from v03_pipeline.lib.tasks.reference_data.write_cached_reference_dataset_query import (
     WriteCachedReferenceDatasetQuery,
 )
@@ -15,7 +14,6 @@
 )
 
 __all__ = [
-    'DeleteOldRunsTask',
     'UpdateProjectTableTask',
     'UpdateLookupTableTask',
     'UpdateVariantAnnotationsTableWithNewSamplesTask',

v03_pipeline/lib/tasks/__init__.py

@@ -13,6 +13,9 @@ def run(self) -> None:
             ht = hl.read_table(self.output().path)
         ht = self.update_table(ht)
         write(ht, self.output().path)
+        # Set force to false after run, allowing "complete()" to succeeded
+        # when dependencies are re-evaluated.
+        self.force = False
 
     def initialize_table(self) -> hl.Table:
         raise NotImplementedError

v03_pipeline/lib/tasks/base/base_update_task.py

@@ -2,7 +2,13 @@
 import luigi
 
 from v03_pipeline.lib.annotations.enums import annotate_enums
-from v03_pipeline.lib.model import Env, ReferenceDatasetCollection
+from v03_pipeline.lib.annotations.rdc_dependencies import (
+    get_rdc_annotation_dependencies,
+)
+from v03_pipeline.lib.model import (
+    Env,
+    ReferenceDatasetCollection,
+)
 from v03_pipeline.lib.paths import (
     valid_reference_dataset_collection_path,
     variant_annotations_table_path,
@@ -17,19 +23,7 @@
 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
+        return get_rdc_annotation_dependencies(self.dataset_type, self.reference_genome)
 
     def output(self) -> luigi.Target:
         return GCSorLocalTarget(

v03_pipeline/lib/tasks/base/base_variant_annotations_table.py

@@ -9,6 +9,9 @@ def run(self) -> None:
         self.init_hail()
         ht = self.create_table()
         write(ht, self.output().path)
+        # Set force to false after run, allowing "complete()" to succeeded
+        # when dependencies are re-evaluated.
+        self.force = False
 
     def create_table(self) -> hl.Table:
         raise NotImplementedError