merge

Author: bpblanken

requirements.txt

@@ -218,7 +218,7 @@ parso==0.8.3
     # via jedi
 pexpect==4.9.0
     # via ipython
-pillow==10.2.0
+pillow==10.3.0
     # via bokeh
 plotly==5.18.0
     # via hail
@@ -252,7 +252,9 @@ pygments==2.17.2
     #   ipython
     #   rich
 pyjwt[crypto]==2.8.0
-    # via msal
+    # via
+    #   msal
+    #   pyjwt
 pyparsing==3.1.1
     # via httplib2
 pyspark==3.3.3

v03_pipeline/bin/write_cached_reference_dataset_query_ht.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+import argparse
+
+import hail as hl
+
+from v03_pipeline.lib.misc.io import write
+from v03_pipeline.lib.model import (
+    CachedReferenceDatasetQuery,
+    DatasetType,
+    ReferenceDatasetCollection,
+    ReferenceGenome,
+)
+from v03_pipeline.lib.paths import (
+    valid_cached_reference_dataset_query_path,
+    valid_reference_dataset_collection_path,
+)
+from v03_pipeline.lib.reference_data.config import CONFIG
+from v03_pipeline.lib.reference_data.dataset_table_operations import (
+    import_ht_from_config_path,
+)
+
+
+def get_ht(
+    dataset_type: DatasetType,
+    reference_genome: ReferenceGenome,
+    query: CachedReferenceDatasetQuery,
+) -> hl.Table:
+    # If the query is defined over an uncombined reference dataset, use the combiner config.
+    if query.query_raw_dataset:
+        config = CONFIG[query.dataset(dataset_type)][reference_genome.v02_value]
+        return import_ht_from_config_path(
+            config,
+            query.dataset(dataset_type),
+            reference_genome,
+        )
+    return hl.read_table(
+        valid_reference_dataset_collection_path(
+            reference_genome,
+            dataset_type,
+            ReferenceDatasetCollection.COMBINED,
+        ),
+    )
+
+
+def run(
+    dataset_type: DatasetType,
+    reference_genome: ReferenceGenome,
+    query: CachedReferenceDatasetQuery,
+):
+    ht = get_ht(dataset_type, reference_genome, query)
+    ht = query.query(ht, dataset_type=dataset_type, reference_genome=reference_genome)
+    destination_path = valid_cached_reference_dataset_query_path(
+        reference_genome,
+        dataset_type,
+        query,
+    )
+    print(f'Uploading ht to {destination_path}')
+    write(ht, destination_path)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--reference-genome',
+        type=ReferenceGenome,
+        choices=list(ReferenceGenome),
+        default=ReferenceGenome.GRCh38,
+    )
+    parser.add_argument(
+        '--dataset-type',
+        type=DatasetType,
+        choices=list(DatasetType),
+        default=None,
+        help='When used, update the passed dataset, otherwise run all datasets.',
+    )
+    parser.add_argument(
+        '--query',
+        type=CachedReferenceDatasetQuery,
+        choices=list(CachedReferenceDatasetQuery),
+        required=True,
+    )
+    args, _ = parser.parse_known_args()
+    if (
+        args.query
+        and args.query
+        not in CachedReferenceDatasetQuery.for_reference_genome_dataset_type(
+            args.reference_genome,
+            args.dataset_type,
+        )
+    ):
+        msg = f'{args.query} is not a valid query for {DatasetType}'
+        raise ValueError(msg)
+    run(args.dataset_type, args.reference_genome, args.query)