lookup sample_type migration (#867)

* add sampletype to family table file path

* all sample types

* missed 1

* project tables

* explicit in family table delete

* use enum

* optional parameter but second DeleteProjectFamilyTablesTask test fails

* oops

* should work

* cleaner

* v3.1 and handle no project tables

* missed one

* paths test

* add lookup migration for sample_type

* old

* no array

* do not modify updates

* comments

* use param_kwargs

* Update base_migration.py

---------

Co-authored-by: Benjamin Blankenmeister <bblanken@broadinstitute.org>

Author: jklugherz

v03_pipeline/lib/migration/base_migration.py

@@ -12,5 +12,5 @@ class BaseMigration(ABC):
 
     @staticmethod
     @abstractmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **kwargs) -> hl.Table:
         pass

v03_pipeline/lib/tasks/base/base_migrate.py

@@ -44,7 +44,7 @@ def update_table(self, ht: hl.Table) -> hl.Table:
         if (
             (self.reference_genome, self.dataset_type)
         ) in migration.reference_genome_dataset_types:
-            ht = migration.migrate(ht)
+            ht = migration.migrate(ht, **self.param_kwargs)
             return ht.annotate_globals(
                 migrations=ht.globals.migrations.append(self.migration_name),
             )

v03_pipeline/lib/tasks/migrate_variant_annotations_table_test.py

@@ -19,7 +19,7 @@ class MockMigration(BaseMigration):
     )
 
     @staticmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **_) -> hl.Table:
         ht = ht.annotate(
             variant_id_id=hl.format('%s_id', ht.variant_id),
         )
@@ -34,7 +34,7 @@ class MockMigration2(BaseMigration):
     )
 
     @staticmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **_) -> hl.Table:
         return ht.annotate_globals(mock_migration2='a second mock migration')
 
 

v03_pipeline/migrations/annotations/0001_add_migrations_global.py

@@ -18,5 +18,5 @@ class AddMigrationsGlobals(BaseMigration):
     )
 
     @staticmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **_) -> hl.Table:
         return ht.annotate_globals(migrations=hl.empty_array(hl.tstr))

v03_pipeline/migrations/annotations/0002_add_remap_pedigree_hash.py

@@ -18,7 +18,7 @@ class AddRemapPedigreeHash(BaseMigration):
     )
 
     @staticmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **_) -> hl.Table:
         return ht.annotate_globals(
             updates=ht.globals.updates.map(
                 lambda u: u.annotate(remap_pedigree_hash=hl.missing(hl.tint32)),

v03_pipeline/migrations/annotations/0003_add_rg38_locus.py

@@ -13,7 +13,7 @@ class AddRG38Locus(BaseMigration):
     )
 
     @staticmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **_) -> hl.Table:
         return ht.annotate(
             rg38_locus=snv_indel.rg38_locus(ht, Env.GRCH37_TO_GRCH38_LIFTOVER_REF_PATH),
         )

v03_pipeline/migrations/lookup/0001_remove_null_families.py

@@ -16,7 +16,7 @@ class RemoveNullFamilies(BaseMigration):
     )
 
     @staticmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **_) -> hl.Table:
         ht = ht.annotate(
             project_stats=ht.project_stats.map(
                 lambda ps: hl.or_missing(hl.any(ps.map(hl.is_defined)), ps),

v03_pipeline/migrations/lookup/0002_add_remap_pedigree_hash.py

@@ -16,7 +16,7 @@ class AddRemapPedigreeHash(BaseMigration):
     )
 
     @staticmethod
-    def migrate(ht: hl.Table) -> hl.Table:
+    def migrate(ht: hl.Table, **_) -> hl.Table:
         return ht.annotate_globals(
             updates=ht.globals.updates.map(
                 lambda u: u.annotate(remap_pedigree_hash=hl.missing(hl.tint32)),

v03_pipeline/migrations/lookup/0003_add_sample_type.py

@@ -0,0 +1,84 @@
+import hail as hl
+import hailtop.fs as hfs
+
+from v03_pipeline.lib.logger import get_logger
+from v03_pipeline.lib.migration.base_migration import BaseMigration
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+from v03_pipeline.lib.paths import project_table_path
+
+logger = get_logger(__name__)
+
+
+class AddLookupSampleType(BaseMigration):
+    reference_genome_dataset_types: frozenset[
+        tuple[ReferenceGenome, DatasetType]
+    ] = frozenset(
+        (
+            (ReferenceGenome.GRCh37, DatasetType.SNV_INDEL),
+            (ReferenceGenome.GRCh38, DatasetType.SNV_INDEL),
+            (ReferenceGenome.GRCh38, DatasetType.MITO),
+        ),
+    )
+
+    @staticmethod
+    def migrate(
+        ht: hl.Table,
+        reference_genome: ReferenceGenome,
+        dataset_type: DatasetType,
+        **_,
+    ) -> hl.Table:
+        """
+        Renames project_guids to project_sample_types.
+        Adds sample_type to global fields project_sample_types and project_families.
+        Assumes that only one project_ht exists for each project across both sample types.
+
+        Old Global fields:
+            'project_guids': array<str>
+            'project_families': dict<str, array<str>>
+            'updates': set<struct {
+                callset: str,
+                project_guid: str,
+                remap_pedigree_hash: int32
+            }>
+        New Global fields:
+            'project_sample_types': array<tuple (
+                str,
+                str
+            )>
+            'project_families': dict<tuple (
+                str,
+                str
+            ), array<str>>
+            'updates': set<struct {
+                callset: str,
+                project_guid: str,
+                remap_pedigree_hash: int32
+            }>
+        """
+        ht = ht.transmute_globals(
+            project_sample_types=ht.globals.project_guids,
+        )
+        collected_globals = ht.globals.collect()[0]
+        project_sample_types = collected_globals['project_sample_types']
+        project_families = collected_globals['project_families']
+
+        for i, project_guid in enumerate(project_sample_types):
+            for sample_type in SampleType:
+                project_ht_path = project_table_path(
+                    reference_genome,
+                    dataset_type,
+                    sample_type,
+                    project_guid,
+                )
+                if not hfs.exists(project_ht_path):
+                    continue
+
+                key = (project_guid, sample_type.value)
+                project_sample_types[i] = key
+                project_families[key] = project_families.pop(project_guid)
+                break
+
+        return ht.annotate_globals(
+            project_sample_types=project_sample_types,
+            project_families=project_families,
+        )