Ensure rows are deleted after deleting samples! (#770)

* Delete project tasks

* cleanup

* ruff format

* well

* rename

* hacking away

* almost there!

* ruff

* Fix missing updates change

* ruff

* Remove debug code

* remove bad merge

* more precision in test

* project table

* allow for missing project

* remove some unnecessary checks

* test already deleted family

* Lots of renames

* More updates

* Sketch

* Flesh out test

* fix paths

* Rename base hail table

* a bunch more renames

* delete project table

* Add delete project families

* add comment

* test it!

* Fix

* add dep

* Lookup table filtering

* Ensure rows with no projects/families defined are removed

* ruff

* remove mock

* Remove mocks from args

* tweak tests

Author: bpblanken

v03_pipeline/lib/misc/lookup.py

@@ -98,15 +98,15 @@ def remove_family_guids(
             )
         ),
     )
-    project_i = ht.project_guids.index(
-        project_guid,
-    )  # double reference because new expression
+    ht = ht.filter(
+        hl.any(ht.project_stats.map(lambda fs: hl.any(fs.map(hl.is_defined)))),
+    )
     return ht.annotate_globals(
         project_families=hl.dict(
-            hl.enumerate(ht.project_families.items()).starmap(
-                lambda i, item: (
+            ht.project_families.items().map(
+                lambda item: (
                     hl.if_else(
-                        i != project_i,
+                        item[0] != project_guid,
                         item,
                         (
                             item[0],
@@ -140,6 +140,7 @@ def remove_project(
             )
         ),
     )
+    ht = ht.filter(hl.any(ht.project_stats.map(hl.is_defined)))
     return ht.annotate_globals(
         project_guids=project_indexes_to_keep.map(
             lambda i: ht.project_guids[i],

v03_pipeline/lib/misc/lookup_test.py

@@ -96,16 +96,8 @@ def test_remove_family_guids(self) -> None:
                     'id': 0,
                     'project_stats': [
                         [
-                            hl.Struct(
-                                ref_samples=0,
-                                heteroplasmic_samples=0,
-                                homoplasmic_samples=0,
-                            ),
-                            hl.Struct(
-                                ref_samples=1,
-                                heteroplasmic_samples=1,
-                                homoplasmic_samples=1,
-                            ),
+                            None,
+                            None,
                             hl.Struct(
                                 ref_samples=2,
                                 heteroplasmic_samples=2,
@@ -125,11 +117,7 @@ def test_remove_family_guids(self) -> None:
                     'id': 1,
                     'project_stats': [
                         [
-                            hl.Struct(
-                                ref_samples=0,
-                                heteroplasmic_samples=0,
-                                homoplasmic_samples=0,
-                            ),
+                            None,
                             hl.Struct(
                                 ref_samples=1,
                                 heteroplasmic_samples=1,
@@ -201,19 +189,6 @@ def test_remove_family_guids(self) -> None:
         self.assertCountEqual(
             lookup_ht.collect(),
             [
-                hl.Struct(
-                    id=0,
-                    project_stats=[
-                        [
-                            hl.Struct(
-                                ref_samples=1,
-                                heteroplasmic_samples=1,
-                                homoplasmic_samples=1,
-                            ),
-                        ],
-                        [],
-                    ],
-                ),
                 hl.Struct(
                     id=1,
                     project_stats=[
@@ -253,13 +228,7 @@ def test_remove_project(self) -> None:
                                 homoplasmic_samples=2,
                             ),
                         ],
-                        [
-                            hl.Struct(
-                                ref_samples=3,
-                                heteroplasmic_samples=3,
-                                homoplasmic_samples=3,
-                            ),
-                        ],
+                        None,
                     ],
                 },
                 {
@@ -330,18 +299,6 @@ def test_remove_project(self) -> None:
         self.assertCountEqual(
             lookup_ht.collect(),
             [
-                hl.Struct(
-                    id=0,
-                    project_stats=[
-                        [
-                            hl.Struct(
-                                ref_samples=3,
-                                heteroplasmic_samples=3,
-                                homoplasmic_samples=3,
-                            ),
-                        ],
-                    ],
-                ),
                 hl.Struct(
                     id=1,
                     project_stats=[

v03_pipeline/lib/tasks/update_lookup_table_with_deleted_families.py

@@ -7,27 +7,12 @@
 from v03_pipeline.lib.tasks.base.base_update_lookup_table import (
     BaseUpdateLookupTableTask,
 )
-from v03_pipeline.lib.tasks.update_variant_annotations_table_with_deleted_families import (
-    UpdateVariantAnnotationsTableWithDeletedFamiliesTask,
-)
 
 
 class UpdateLookupTableWithDeletedFamiliesTask(BaseUpdateLookupTableTask):
     project_guid = luigi.Parameter()
     family_guids = luigi.ListParameter()
 
-    def requires(self) -> luigi.Task:
-        # We require updating the annotations table first so that
-        # we are able to use the lookup table to determine which rows
-        # of the annotations table require re-annotation.
-        return UpdateVariantAnnotationsTableWithDeletedFamiliesTask(
-            dataset_type=self.dataset_type,
-            sample_type=self.sample_type,
-            reference_genome=self.reference_genome,
-            project_guid=self.project_guid,
-            family_guids=self.family_guids,
-        )
-
     def complete(self) -> bool:
         return super().complete() and hl.eval(
             hl.bind(

v03_pipeline/lib/tasks/update_lookup_table_with_deleted_families_test.py

@@ -7,19 +7,13 @@
 from v03_pipeline.lib.tasks.update_lookup_table_with_deleted_families import (
     UpdateLookupTableWithDeletedFamiliesTask,
 )
-from v03_pipeline.lib.test.mock_complete_task import MockCompleteTask
 from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
 
 
-@mock.patch(
-    'v03_pipeline.lib.tasks.update_lookup_table_with_deleted_families.UpdateVariantAnnotationsTableWithDeletedFamiliesTask',
-)
 class UpdateLookupTableWithDeletedProjectTaskTest(MockedDatarootTestCase):
     def test_delete_project_empty_table(
         self,
-        mock_update_lookup_table_task: mock.Mock,
     ) -> None:
-        mock_update_lookup_table_task.return_value = MockCompleteTask()
         worker = luigi.worker.Worker()
         task = UpdateLookupTableWithDeletedFamiliesTask(
             dataset_type=DatasetType.SNV_INDEL,
@@ -51,9 +45,7 @@ def test_delete_project_empty_table(
     def test_delete_project(
         self,
         mock_initialize_table: mock.Mock,
-        mock_update_lookup_table_task: mock.Mock,
     ) -> None:
-        mock_update_lookup_table_task.return_value = MockCompleteTask()
         mock_initialize_table.return_value = hl.Table.parallelize(
             [
                 {

v03_pipeline/lib/tasks/update_lookup_table_with_deleted_project.py

@@ -7,22 +7,11 @@
 from v03_pipeline.lib.tasks.base.base_update_lookup_table import (
     BaseUpdateLookupTableTask,
 )
-from v03_pipeline.lib.tasks.update_variant_annotations_table_with_deleted_project import (
-    UpdateVariantAnnotationsTableWithDeletedProjectTask,
-)
 
 
 class UpdateLookupTableWithDeletedProjectTask(BaseUpdateLookupTableTask):
     project_guid = luigi.Parameter()
 
-    def requires(self) -> luigi.Task:
-        return UpdateVariantAnnotationsTableWithDeletedProjectTask(
-            dataset_type=self.dataset_type,
-            sample_type=self.sample_type,
-            reference_genome=self.reference_genome,
-            project_guid=self.project_guid,
-        )
-
     def complete(self) -> bool:
         return super().complete() and hl.eval(
             ~hl.read_table(self.output().path).updates.project_guid.contains(

v03_pipeline/lib/tasks/update_lookup_table_with_deleted_project_test.py

@@ -7,19 +7,13 @@
 from v03_pipeline.lib.tasks.update_lookup_table_with_deleted_project import (
     UpdateLookupTableWithDeletedProjectTask,
 )
-from v03_pipeline.lib.test.mock_complete_task import MockCompleteTask
 from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
 
 
-@mock.patch(
-    'v03_pipeline.lib.tasks.update_lookup_table_with_deleted_project.UpdateVariantAnnotationsTableWithDeletedProjectTask',
-)
 class UpdateLookupTableWithDeletedProjectTaskTest(MockedDatarootTestCase):
     def test_delete_project_empty_table(
         self,
-        mock_update_lookup_table_task: mock.Mock,
     ) -> None:
-        mock_update_lookup_table_task.return_value = MockCompleteTask()
         worker = luigi.worker.Worker()
         task = UpdateLookupTableWithDeletedProjectTask(
             dataset_type=DatasetType.SNV_INDEL,
@@ -50,9 +44,7 @@ def test_delete_project_empty_table(
     def test_delete_project(
         self,
         mock_initialize_table: mock.Mock,
-        mock_update_lookup_table_task: mock.Mock,
     ) -> None:
-        mock_update_lookup_table_task.return_value = MockCompleteTask()
         mock_initialize_table.return_value = hl.Table.parallelize(
             [
                 {

v03_pipeline/lib/tasks/update_variant_annotations_table_with_deleted_families.py

@@ -2,16 +2,12 @@
 import luigi
 
 from v03_pipeline.lib.annotations.fields import get_fields
-from v03_pipeline.lib.misc.lookup import (
-    remove_family_guids,
-)
-from v03_pipeline.lib.paths import (
-    lookup_table_path,
-)
 from v03_pipeline.lib.tasks.base.base_update_variant_annotations_table import (
     BaseUpdateVariantAnnotationsTableTask,
 )
-from v03_pipeline.lib.tasks.files import HailTableTask
+from v03_pipeline.lib.tasks.update_lookup_table_with_deleted_families import (
+    UpdateLookupTableWithDeletedFamiliesTask,
+)
 
 
 class UpdateVariantAnnotationsTableWithDeletedFamiliesTask(
@@ -25,72 +21,33 @@ def __init__(self, *args, **kwargs):
         self.done = False
 
     def requires(self) -> luigi.Task:
-        return HailTableTask(
-            lookup_table_path(
-                self.reference_genome,
-                self.dataset_type,
-            ),
+        return UpdateLookupTableWithDeletedFamiliesTask(
+            dataset_type=self.dataset_type,
+            sample_type=self.sample_type,
+            reference_genome=self.reference_genome,
+            project_guid=self.project_guid,
+            family_guids=self.family_guids,
         )
 
     def complete(self) -> bool:
         if not self.dataset_type.has_lookup_table:
             return True
-        return super().complete() and (
-            # We don't have the concept of families being present or
-            # not present in annotations table, thus we check the "lookup"
-            # table OR a done flag to prevent the task from looping over itself.
-            self.done
-            or hl.eval(
-                hl.bind(
-                    lambda family_guids: (
-                        hl.is_missing(family_guids)  # The project itself is missing
-                        | hl.all(
-                            hl.array(list(self.family_guids)).map(
-                                lambda family_guid: ~hl.set(family_guids).contains(
-                                    family_guid,
-                                ),
-                            ),
-                        )
-                    ),
-                    hl.read_table(self.input().path).globals.project_families.get(
-                        self.project_guid,
-                    ),
-                ),
-            )
-        )
+        # We don't have the concept of families being present or not present
+        #  in annotations table a done flag to prevent the task from looping over itself.
+        return super().complete() and self.done
 
     def update_table(self, ht: hl.Table) -> hl.Table:
         if not self.dataset_type.has_lookup_table:
             return ht
         lookup_ht = hl.read_table(self.input().path)
-        project_i = hl.eval(lookup_ht.globals.project_guids.index(self.project_guid))
-
-        # Filter lookup table to only the rows where any of the requested families are defined
-        # (A family may be set to missing during lookup table construction if all values for a family are 0)
-        family_guids_set = hl.set(list(self.family_guids))
-        family_indexes = (
-            hl.enumerate(lookup_ht.globals.project_families[self.project_guid])
-            .filter(lambda item: family_guids_set.contains(item[1]))
-            .map(lambda item: item[0])
-        )
-        lookup_ht = lookup_ht.filter(
-            hl.any(
-                family_indexes.map(
-                    lambda i: hl.is_defined(lookup_ht.project_stats[project_i][i]),
-                ),
-            ),
-        )
-        lookup_ht = remove_family_guids(lookup_ht, self.project_guid, family_guids_set)
-        family_variants_ht = ht.semi_join(lookup_ht)
-        family_variants_ht = family_variants_ht.annotate(
+        ht = ht.semi_join(lookup_ht)
+        ht = ht.annotate(
             **get_fields(
-                family_variants_ht,
+                ht,
                 self.dataset_type.lookup_table_annotation_fns,
                 lookup_ht=lookup_ht,
                 **self.param_kwargs,
             ),
         )
-        ht = ht.anti_join(lookup_ht)
-        ht = ht.union(family_variants_ht)
         self.done = True
         return ht

v03_pipeline/lib/tasks/update_variant_annotations_table_with_deleted_families_test.py

@@ -84,6 +84,10 @@ def setUp(self) -> None:
             globals=hl.Struct(
                 project_guids=['project_a', 'project_b'],
                 project_families={'project_a': ['1', '2', '3'], 'project_b': ['4']},
+                updates={
+                    hl.Struct(callset='abc', project_guid='project_a'),
+                    hl.Struct(callset='123', project_guid='project_b'),
+                },
             ),
         )
         ht.write(
@@ -158,10 +162,6 @@ def test_update_annotations_with_deleted_project(self) -> None:
             ht.collect(),
             [
                 hl.Struct(id=0, gt_stats=hl.Struct(AC=9, AN=18, AF=0.5, hom=3)),
-                hl.Struct(id=1, gt_stats=hl.Struct(AC=0, AN=1, AF=0.25, hom=0)),
-                hl.Struct(
-                    id=2,
-                    gt_stats=hl.Struct(AC=0, AN=1, AF=0.25, hom=0),
-                ),
+                hl.Struct(id=1, gt_stats=hl.Struct(AC=12, AN=24, AF=0.5, hom=4)),
             ],
         )