feat: add tasks for writing parquet tables. (#1076)

* Formatting function enums

* improve function name

* handle lookups that propagate missing

* format

* ruff

* rename func

* second func

* start test

* finish tests

* ruff

* Add camelcase

* add camelcase

* ruff

* missing init py

* move functions to export

* mostly functioning entries task

* import

* improve private reference datasets logic

* first pass

* progress

* tests passing

* key

* annotations table

* test transcripts

* closer

* missed some tests

* add a parquet reader requirement

* use project tables

* finish off first pass at test!

* reformat test

* another batch

* add filters

* ruff

* handle hgmd edge case

* add hgmd

* sort it

* ruff

* bugfixes

* update test to new format

* ruff

* v03

* merge

* no longer used

* lint

* formatting

* special case the export

* remove gene/map

* ruff

* add new annotations

* print

* ruff

* canonical is not a float

* add new clinvar variants parquet export (#1092)

* feat: grch37 SNV_INDEL export (#1095)

* support grch37

* enable grch37

* feat: alphabetize nested field (#1096)

* Support alphabetization of nested field

* better sorting

* Update misc_test.py

* ruff

Author: bpblanken

requirements-dev.in

@@ -7,3 +7,4 @@ responses>=0.23.1
 ruff>=0.1.8
 shellcheck-py>=0.10.0
 pysam
+pyarrow

requirements-dev.txt

@@ -59,6 +59,8 @@ packaging==24.2
     #   sphinx
 pip-tools==7.4.1
     # via -r requirements-dev.in
+pyarrow==19.0.1
+    # via -r requirements-dev.in
 pygments==2.19.1
     # via
     #   -c requirements.txt

v03_pipeline/lib/model/dataset_type.py

@@ -1,3 +1,4 @@
+from collections import OrderedDict
 from collections.abc import Callable
 from enum import StrEnum
 
@@ -390,6 +391,10 @@ def filter_invalid_sites(self):
     def should_export_to_vcf(self):
         return self == DatasetType.SV
 
+    @property
+    def should_export_to_parquet(self):
+        return self == DatasetType.SNV_INDEL
+
     @property
     def export_vcf_annotation_fns(self) -> list[Callable[..., hl.Expression]]:
         return {
@@ -400,6 +405,37 @@ def export_vcf_annotation_fns(self) -> list[Callable[..., hl.Expression]]:
             ],
         }[self]
 
+    def export_parquet_filterable_transcripts_fields(
+        self,
+        reference_genome: ReferenceGenome,
+    ) -> OrderedDict[str, str]:
+        fields = ['geneId']
+        if self in {DatasetType.SV, DatasetType.GCNV}:
+            fields = [
+                *fields,
+                'majorConsequence',
+            ]
+        if self in {DatasetType.SNV_INDEL, DatasetType.MITO}:
+            fields = [
+                *fields,
+                'canonical',
+                'consequenceTerms',
+            ]
+        fields = {
+            # above fields are renamed to themselves
+            k: k
+            for k in fields
+        }
+        if self == DatasetType.SNV_INDEL and reference_genome == ReferenceGenome.GRCh38:
+            fields = {
+                **fields,
+                'alphamissensePathogenicity': 'alphamissense.pathogenicity',
+                'extendedIntronicSpliceRegionVariant': 'spliceregion.extended_intronic_splice_region_variant',
+                'fiveutrConsequence': 'utrannotator.fiveutrConsequence',
+            }
+        # Parquet export expects all fields sorted alphabetically
+        return OrderedDict(sorted(fields.items()))
+
     @property
     def overwrite_male_non_par_calls(self) -> None:
         return self == DatasetType.SV

v03_pipeline/lib/model/feature_flag.py

@@ -7,6 +7,7 @@
 )
 CHECK_SEX_AND_RELATEDNESS = os.environ.get('CHECK_SEX_AND_RELATEDNESS') == '1'
 EXPECT_TDR_METRICS = os.environ.get('EXPECT_TDR_METRICS') == '1'
+EXPORT_TO_PARQUET = os.environ.get('EXPORT_TO_PARQUET') == '1'
 INCLUDE_PIPELINE_VERSION_IN_PREFIX = (
     os.environ.get('INCLUDE_PIPELINE_VERSION_IN_PREFIX') == '1'
 )
@@ -21,6 +22,7 @@ class FeatureFlag:
     ACCESS_PRIVATE_REFERENCE_DATASETS: bool = ACCESS_PRIVATE_REFERENCE_DATASETS
     CHECK_SEX_AND_RELATEDNESS: bool = CHECK_SEX_AND_RELATEDNESS
     EXPECT_TDR_METRICS: bool = EXPECT_TDR_METRICS
+    EXPORT_TO_PARQUET: bool = EXPORT_TO_PARQUET
     INCLUDE_PIPELINE_VERSION_IN_PREFIX: bool = INCLUDE_PIPELINE_VERSION_IN_PREFIX
     RUN_PIPELINE_ON_DATAPROC: bool = RUN_PIPELINE_ON_DATAPROC
     SHOULD_TRIGGER_HAIL_BACKEND_RELOAD: bool = SHOULD_TRIGGER_HAIL_BACKEND_RELOAD

v03_pipeline/lib/paths.py

@@ -363,6 +363,66 @@ def variant_annotations_vcf_path(
     )
 
 
+def new_clinvar_variants_parquet_path(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    run_id: str,
+) -> str:
+    return os.path.join(
+        runs_path(
+            reference_genome,
+            dataset_type,
+        ),
+        run_id,
+        'new_clinvar_variants.parquet',
+    )
+
+
+def new_entries_parquet_path(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    run_id: str,
+) -> str:
+    return os.path.join(
+        runs_path(
+            reference_genome,
+            dataset_type,
+        ),
+        run_id,
+        'new_entries.parquet',
+    )
+
+
+def new_transcripts_parquet_path(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    run_id: str,
+) -> str:
+    return os.path.join(
+        runs_path(
+            reference_genome,
+            dataset_type,
+        ),
+        run_id,
+        'new_transcripts.parquet',
+    )
+
+
+def new_variants_parquet_path(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    run_id: str,
+) -> str:
+    return os.path.join(
+        runs_path(
+            reference_genome,
+            dataset_type,
+        ),
+        run_id,
+        'new_variants.parquet',
+    )
+
+
 def new_variants_table_path(
     reference_genome: ReferenceGenome,
     dataset_type: DatasetType,

v03_pipeline/lib/tasks/base/base_write_parquet.py

@@ -0,0 +1,19 @@
+import luigi
+
+from v03_pipeline.lib.misc.io import checkpoint
+from v03_pipeline.lib.tasks.files import GCSorLocalFolderTarget
+
+
+class BaseWriteParquetTask(luigi.Task):
+    def complete(self) -> luigi.Target:
+        return GCSorLocalFolderTarget(self.output().path).exists()
+
+    def run(self) -> None:
+        ht = self.create_table()
+        ht, _ = checkpoint(ht)
+        df = ht.to_spark(flatten=False)
+        df = df.withColumnRenamed('key_', 'key')
+        df.write.parquet(
+            self.output().path,
+            mode='overwrite',
+        )