Cleanup hail

Author: bpblanken

v03_pipeline/lib/misc/family_entries.py

@@ -95,9 +95,8 @@ def deglobalize_ids(ht: hl.Table) -> hl.Table:
 
 def remove_new_callset_family_guids(
     ht: hl.Table,
-    family_guids: list[str],
+    family_guids: hl.SetExpression,
 ) -> hl.Table:
-    family_guids = hl.set(family_guids)
     # Remove families from the existing project table structure (both the entries arrays and the globals are mutated)
     family_indexes_to_keep = hl.array(
         hl.enumerate(ht.globals.family_guids)

v03_pipeline/lib/misc/family_entries_test.py

@@ -244,7 +244,7 @@ def test_remove_new_callset_family_guids(self) -> None:
                 },
             ),
         )
-        family_entries_ht = remove_new_callset_family_guids(family_entries_ht, ['012'])
+        family_entries_ht = remove_new_callset_family_guids(family_entries_ht, hl.set(['012']))
         self.assertCountEqual(
             family_entries_ht.globals.collect(),
             [
@@ -326,7 +326,7 @@ def test_remove_new_callset_family_guids_all_families(self) -> None:
                 },
             ),
         )
-        ht = remove_new_callset_family_guids(family_entries_ht, ['012', '123'])
+        ht = remove_new_callset_family_guids(family_entries_ht, hl.set(['012', '123']))
         self.assertCountEqual(
             ht.globals.collect(),
             [hl.Struct(family_guids=[], family_samples={})],

v03_pipeline/lib/misc/lookup.py

@@ -62,11 +62,10 @@ def globalize_ids(ht: hl.Table, project_guid: str) -> hl.Table:
 def remove_new_callset_family_guids(
     ht: hl.Table,
     project_guid: str,
-    family_guids: list[str],
+    family_guids: hl.SetExpression,
 ) -> hl.Table:
     if project_guid not in hl.eval(ht.globals.project_families):
         return ht
-    family_guids = hl.set(family_guids)
     project_i = ht.project_guids.index(project_guid)
     family_indexes_to_keep = hl.array(
         hl.enumerate(ht.globals.project_families[project_guid])

v03_pipeline/lib/misc/lookup_test.py

@@ -171,17 +171,17 @@ def test_remove_new_callset_family_guids(self) -> None:
         lookup_ht = remove_new_callset_family_guids(
             lookup_ht,
             'project_c',
-            ['2'],
+            hl.set(['2']),
         )
         lookup_ht = remove_new_callset_family_guids(
             lookup_ht,
             'project_a',
-            ['3', '1'],
+            hl.set(['3', '1']),
         )
         lookup_ht = remove_new_callset_family_guids(
             lookup_ht,
             'project_b',
-            ['4'],
+            hl.set(['4']),
         )
         self.assertCountEqual(
             lookup_ht.globals.collect(),

v03_pipeline/lib/tasks/update_lookup_table.py

@@ -104,7 +104,7 @@ def update_table(self, ht: hl.Table) -> hl.Table:
         ht = remove_new_callset_family_guids(
             ht,
             self.project_guid,
-            list(callset_mt.family_samples.collect()[0].keys()),
+            callset_mt.family_samples.key_set(),
         )
         callset_ht = compute_callset_lookup_ht(
             self.dataset_type,

v03_pipeline/lib/tasks/update_project_table.py

@@ -106,7 +106,7 @@ def update_table(self, ht: hl.Table) -> hl.Table:
             )
         ht = remove_new_callset_family_guids(
             ht,
-            list(callset_mt.family_samples.collect()[0].keys()),
+            callset_mt.family_samples.key_set(),
         )
         ht = join_family_entries_hts(ht, callset_ht)
         return ht.select_globals(