refactor: improve code sharing for clickhouse migration tasks

Author: bpblanken

v03_pipeline/lib/misc/family_entries.py

@@ -100,15 +100,17 @@ def globalize_ids(ht: hl.Table) -> hl.Table:
     )
 
 
-def deglobalize_ids(ht: hl.Table) -> hl.Table:
+def deglobalize_ids(ht: hl.Table, sample_id_field: str = 's') -> hl.Table:
     ht = ht.annotate(
         family_entries=(
             hl.enumerate(ht.family_entries).starmap(
                 lambda i, fe: hl.enumerate(fe).starmap(
                     lambda j, e: hl.Struct(
                         **e,
-                        s=ht.family_samples[ht.family_guids[i]][j],
-                        family_guid=ht.family_guids[i],
+                        **{
+                            sample_id_field: ht.family_samples[ht.family_guids[i]][j],
+                            'family_guid': ht.family_guids[i],
+                        },
                     ),
                 ),
             )

v03_pipeline/lib/model/dataset_type.py

@@ -101,6 +101,24 @@ def entries_fields(
             },
         }[self]
 
+    @property
+    def entries_export_fields(
+        self,
+    ) -> Callable[hl.StructExpression, hl.StructExpression]:
+        return {
+            DatasetType.SNV_INDEL: lambda fe: hl.Struct(
+                sampleId=fe.sampleId,
+                gt=hl.case()
+                .when(fe.GT.is_hom_ref(), 0)
+                .when(fe.GT.is_het(), 1)
+                .when(fe.GT.is_hom_var(), 2)
+                .default(hl.missing(hl.tint32)),
+                gq=fe.GQ,
+                ab=fe.AB,
+                dp=fe.DP,
+            ),
+        }[self]
+
     @property
     def row_fields(
         self,

v03_pipeline/lib/tasks/exports/write_new_entries_parquet.py

@@ -6,6 +6,7 @@
 from v03_pipeline.lib.annotations.fields import get_fields
 from v03_pipeline.lib.misc.family_entries import (
     compute_callset_family_entries_ht,
+    deglobalize_ids,
 )
 from v03_pipeline.lib.paths import (
     new_entries_parquet_path,
@@ -90,16 +91,7 @@ def create_table(self) -> None:
                     **self.param_kwargs,
                 ),
             )
-            ht = ht.annotate(
-                family_entries=hl.enumerate(ht.family_entries).starmap(
-                    lambda i, fs: hl.enumerate(fs).starmap(
-                        lambda j, e: e.annotate(
-                            family_guid=ht.family_guids[i],  # noqa: B023
-                            sampleId=ht.family_samples[ht.family_guids[i]][j],  # noqa: B023
-                        ),
-                    ),
-                ),
-            )
+            ht = deglobalize_ids(ht)
             annotations_ht = hl.read_table(
                 self.input()[ANNOTATIONS_TABLE_TASK].path,
             )
@@ -124,20 +116,16 @@ def create_table(self) -> None:
                 family_guid=ht.family_entries.family_guid[0],
                 sample_type=self.sample_type.value,
                 xpos=get_expr_for_xpos(ht.locus),
-                is_gnomad_gt_5_percent=hl.is_defined(ht.is_gt_5_percent),
+                **(
+                    {
+                        'is_gnomad_gt_5_percent': hl.is_defined(ht.is_gt_5_percent),
+                    }
+                    if hasattr(ht, 'is_gnomad_gt_5_percent')
+                    else {}
+                ),
                 filters=ht.filters,
                 calls=ht.family_entries.map(
-                    lambda fe: hl.Struct(
-                        sampleId=fe.sampleId,
-                        gt=hl.case()
-                        .when(fe.GT.is_hom_ref(), 0)
-                        .when(fe.GT.is_het(), 1)
-                        .when(fe.GT.is_hom_var(), 2)
-                        .default(hl.missing(hl.tint32)),
-                        gq=fe.GQ,
-                        ab=fe.AB,
-                        dp=fe.DP,
-                    ),
+                    lambda fe: self.dataset_type.entries_export_fields(fe),
                 ),
                 sign=1,
             )