bugfix: mito annotations incorporate variant fields

Author: bpblanken

v03_pipeline/lib/tasks/exports/fields.py

@@ -1,6 +1,7 @@
 import hail as hl
 
 from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+from v03_pipeline.lib.tasks.exports.misc import reformat_transcripts_for_export
 
 
 def reference_independent_contig(locus: hl.LocusExpression):
@@ -323,7 +324,11 @@ def get_consequences_fields(
             'sortedTranscriptConsequences': ht.sortedTranscriptConsequences,
         },
         DatasetType.MITO: lambda ht: {
-            'sortedTranscriptConsequences': ht.sortedTranscriptConsequences,
+            # MITO transcripts are not exported to their own table,
+            # but the structure should be preserved here.
+            'sortedTranscriptConsequences': hl.enumerate(
+                ht.sortedTranscriptConsequences,
+            ).starmap(reformat_transcripts_for_export),
         },
         DatasetType.SV: lambda ht: {
             'sortedGeneConsequences': ht.sortedGeneConsequences,

v03_pipeline/lib/tasks/exports/misc.py

@@ -47,7 +47,7 @@ def array_structexpression_fields(ht: hl.Table):
 
 
 def reformat_transcripts_for_export(i: int, s: hl.StructExpression):
-    return (
+    formatted_s = (
         s.annotate(
             majorConsequence=s.consequenceTerms.first(),
             transcriptRank=i,
@@ -62,6 +62,7 @@ def reformat_transcripts_for_export(i: int, s: hl.StructExpression):
             transcriptRank=i,
         ).drop('isLofNagnag', 'lofFilters')
     )
+    return sorted_hl_struct(formatted_s)
 
 
 def export_parquet_filterable_transcripts_fields(

v03_pipeline/lib/tasks/exports/write_new_transcripts_parquet.py

@@ -13,7 +13,6 @@
 from v03_pipeline.lib.tasks.exports.misc import (
     camelcase_array_structexpression_fields,
     reformat_transcripts_for_export,
-    sorted_hl_struct,
     unmap_formatting_annotation_enums,
 )
 from v03_pipeline.lib.tasks.files import GCSorLocalFolderTarget, GCSorLocalTarget
@@ -62,7 +61,5 @@ def create_table(self) -> None:
             key_=ht.key_,
             transcripts=hl.enumerate(
                 ht.sortedTranscriptConsequences,
-            )
-            .starmap(reformat_transcripts_for_export)
-            .map(sorted_hl_struct),
+            ).starmap(reformat_transcripts_for_export),
         )

v03_pipeline/lib/tasks/exports/write_new_variants_parquet_test.py

@@ -407,16 +407,20 @@ def test_mito_write_new_variants_parquet_test(
                     'sortedTranscriptConsequences': [
                         {
                             'aminoAcids': None,
+                            'biotype': 'Mt_tRNA',
                             'canonical': 1,
                             'codons': None,
+                            'consequenceTerms': ['non_coding_transcript_exon_variant'],
                             'geneId': 'ENSG00000210049',
                             'hgvsc': 'ENST00000387314.1:n.2T>C',
                             'hgvsp': None,
+                            'loftee': {
+                                'isLofNagnag': None,
+                                'lofFilters': None,
+                            },
+                            'majorConsequence': 'non_coding_transcript_exon_variant',
                             'transcriptId': 'ENST00000387314',
-                            'isLofNagnag': None,
-                            'biotype': 'Mt_tRNA',
-                            'consequenceTerms': ['non_coding_transcript_exon_variant'],
-                            'lofFilters': None,
+                            'transcriptRank': 0,
                         },
                     ],
                 },