Merge pull request #838 from broadinstitute/benb/intersect_family_guids

Intersect family guids when determining which to load

Author: bpblanken

v03_pipeline/lib/tasks/write_project_family_tables.py

@@ -2,8 +2,11 @@
 import luigi
 import luigi.util
 
+from v03_pipeline.lib.misc.io import import_pedigree
+from v03_pipeline.lib.misc.pedigree import parse_pedigree_ht_to_families
 from v03_pipeline.lib.tasks.base.base_hail_table import BaseHailTableTask
 from v03_pipeline.lib.tasks.base.base_loading_run_params import BaseLoadingRunParams
+from v03_pipeline.lib.tasks.files import RawFileTask
 from v03_pipeline.lib.tasks.update_project_table import UpdateProjectTableTask
 from v03_pipeline.lib.tasks.write_family_table import WriteFamilyTableTask
 
@@ -30,13 +33,26 @@ def complete(self) -> bool:
 
     def run(self):
         # https://luigi.readthedocs.io/en/stable/tasks.html#dynamic-dependencies
+        # Fetch family guids from project table
         update_project_table_task: luigi.Target = yield self.clone(
             UpdateProjectTableTask,
             force=False,
         )
         project_ht = hl.read_table(update_project_table_task.path)
-        family_guids = hl.eval(project_ht.globals.family_guids)
-        for family_guid in family_guids:
+        family_guids_in_project_table = set(hl.eval(project_ht.globals.family_guids))
+
+        # Fetch family guids from pedigree
+        pedigree_ht_task: luigi.Target = yield RawFileTask(self.project_pedigree_path)
+        pedigree_ht = import_pedigree(pedigree_ht_task.path)
+        families_guids_in_pedigree = {
+            f.family_guid for f in parse_pedigree_ht_to_families(pedigree_ht)
+        }
+
+        # Intersect them
+        family_guids_to_load = (
+            family_guids_in_project_table & families_guids_in_pedigree
+        )
+        for family_guid in family_guids_to_load:
             self.dynamic_write_family_table_tasks.add(
                 self.clone(WriteFamilyTableTask, family_guid=family_guid),
             )

v03_pipeline/lib/tasks/write_project_family_tables_test.py

@@ -2,6 +2,7 @@
 import luigi.worker
 
 from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+from v03_pipeline.lib.paths import project_table_path
 from v03_pipeline.lib.tasks.write_project_family_tables import (
     WriteProjectFamilyTablesTask,
 )
@@ -10,9 +11,12 @@
 TEST_SNV_INDEL_VCF = 'v03_pipeline/var/test/callsets/1kg_30variants.vcf'
 TEST_REMAP = 'v03_pipeline/var/test/remaps/test_remap_1.tsv'
 TEST_PEDIGREE_4 = 'v03_pipeline/var/test/pedigrees/test_pedigree_4.tsv'
+TEST_PEDIGREE_4_SUBSET = 'v03_pipeline/var/test/pedigrees/test_pedigree_4_subset.tsv'
 
 
 class WriteProjectFamilyTablesTest(MockedDatarootTestCase):
+    maxDiff = None
+
     def test_snv_write_project_family_tables_task(self) -> None:
         worker = luigi.worker.Worker()
         write_project_family_tables = WriteProjectFamilyTablesTask(
@@ -51,3 +55,54 @@ def test_snv_write_project_family_tables_task(self) -> None:
                 [['NA20888_1']],
             ],
         )
+
+        write_project_family_tables_subset = WriteProjectFamilyTablesTask(
+            reference_genome=ReferenceGenome.GRCh38,
+            dataset_type=DatasetType.SNV_INDEL,
+            sample_type=SampleType.WGS,
+            callset_path=TEST_SNV_INDEL_VCF,
+            project_guid='R0113_test_project',
+            project_remap_path=TEST_REMAP,
+            project_pedigree_path=TEST_PEDIGREE_4_SUBSET,
+            skip_validation=True,
+            skip_check_sex_and_relatedness=True,
+        )
+        worker.add(write_project_family_tables_subset)
+        worker.run()
+        self.assertTrue(write_project_family_tables_subset.complete())
+        hts = [
+            hl.read_table(write_family_table_task.output().path)
+            for write_family_table_task in write_project_family_tables_subset.dynamic_write_family_table_tasks
+        ]
+        # Only one family table written
+        self.assertEqual(
+            len(hts),
+            1,
+        )
+        # Project table still contains all family guids
+        self.assertCountEqual(
+            hl.read_table(
+                project_table_path(
+                    ReferenceGenome.GRCh38,
+                    DatasetType.SNV_INDEL,
+                    'R0113_test_project',
+                ),
+            ).family_guids.collect(),
+            [
+                [
+                    '123_1',
+                    '234_1',
+                    '345_1',
+                    '456_1',
+                    '567_1',
+                    '678_1',
+                    '789_1',
+                    '890_1',
+                    '901_1',
+                    'bcd_1',
+                    'cde_1',
+                    'def_1',
+                    'efg_1',
+                ],
+            ],
+        )

v03_pipeline/var/test/pedigrees/test_pedigree_4_subset.tsv

@@ -0,0 +1,2 @@
+Project_GUID	Family_GUID	Family_ID	Individual_ID	Paternal_ID	Maternal_ID	Sex
+R0114_project4	123_1	123	NA19675_1			F