create updated reference dataset collection task (#656)

* first pass at update task

* move stuff back to /lib file

* add tests for former combine file

* add tests for former combine file

* comment w/ logic for complete methjod

* task unit test

* add missing test cases and default luigi datasets param

* delete_reference_data_ht removed

* remove test setup tmpdir stuff

Author: jklugherz

download_and_create_reference_datasets/v02/hail_scripts/write_combined_interval_ref_data.py

@@ -3,7 +3,7 @@
 
 import hail as hl
 
-from v03_pipeline.lib.reference_data.combine import join_hts
+from v03_pipeline.lib.reference_data.dataset_table_operations import join_hts
 
 VERSION = '2.0.5'
 OUTPUT_PATH = "gs://seqr-reference-data/GRCh38/combined_interval_reference_data/combined_interval_reference_data.ht"

download_and_create_reference_datasets/v02/hail_scripts/write_combined_reference_data_ht.py

@@ -3,7 +3,7 @@
 
 import hail as hl
 
-from v03_pipeline.lib.reference_data.combine import join_hts
+from v03_pipeline.lib.reference_data.dataset_table_operations import join_hts
 from v03_pipeline.lib.reference_data.config import CONFIG
 
 VERSION = '2.0.4'

download_and_create_reference_datasets/v02/mito/write_combined_mito_reference_data_hts.py

@@ -4,7 +4,7 @@
 
 import hail as hl
 
-from v03_pipeline.lib.reference_data.combine import join_hts
+from v03_pipeline.lib.reference_data.dataset_table_operations import join_hts
 
 VERSION = '2.0.4'
 OUTPUT_PATH = 'gs://seqr-reference-data/GRCh38/mitochondrial/all_mito_reference_data/combined_reference_data_chrM.ht'

luigi_pipeline/tests/test_hail_tasks.py

@@ -40,12 +40,10 @@ def _hail_matrix_table_task(self):
 
     def _set_validation_configs(self):
         global_config = GlobalConfig()
-        global_config.param_kwargs[
-            'validation_37_coding_ht'
-        ] = global_config.validation_37_coding_ht = 'tests/data/validation_37_coding.ht'
-        global_config.param_kwargs[
-            'validation_37_noncoding_ht'
-        ] = (
+        global_config.param_kwargs['validation_37_coding_ht'] = (
+            global_config.validation_37_coding_ht
+        ) = 'tests/data/validation_37_coding.ht'
+        global_config.param_kwargs['validation_37_noncoding_ht'] = (
             global_config.validation_37_noncoding_ht
         ) = 'tests/data/validation_37_noncoding.ht'
 

pyproject.toml

@@ -67,3 +67,6 @@ inline-quotes = "single"
 
 [tool.ruff.pylint]
 max-args = 6
+
+[tool.ruff.format]
+quote-style = "single"

v03_pipeline/bin/write_combined_interval_reference_ht.py

@@ -13,7 +13,10 @@
     ReferenceGenome,
 )
 from v03_pipeline.lib.paths import valid_reference_dataset_collection_path
-from v03_pipeline.lib.reference_data.combine import join_hts, update_existing_joined_hts
+from v03_pipeline.lib.reference_data.dataset_table_operations import (
+    join_hts,
+    update_existing_joined_hts,
+)
 
 
 def run(dataset_type: DatasetType, dataset: str | None):

v03_pipeline/bin/write_combined_reference_ht.py

@@ -14,7 +14,10 @@
     ReferenceGenome,
 )
 from v03_pipeline.lib.paths import valid_reference_dataset_collection_path
-from v03_pipeline.lib.reference_data.combine import join_hts, update_existing_joined_hts
+from v03_pipeline.lib.reference_data.dataset_table_operations import (
+    join_hts,
+    update_existing_joined_hts,
+)
 
 
 def run(

v03_pipeline/bin/write_hgmd_ht.py

@@ -8,7 +8,7 @@
     ReferenceGenome,
 )
 from v03_pipeline.lib.paths import valid_reference_dataset_collection_path
-from v03_pipeline.lib.reference_data.combine import join_hts
+from v03_pipeline.lib.reference_data.dataset_table_operations import join_hts
 
 
 def run(reference_genome: ReferenceGenome):

v03_pipeline/lib/misc/sample_entries.py

@@ -8,9 +8,7 @@ def globalize_sample_ids(ht: hl.Table) -> hl.Table:
             # NB: normal python expression here because the row is localized.
             # I had an easier time with this than hl.agg.take(1), which was an
             # alternative implementation.
-            [e.s for e in row[0].entries]
-            if len(row) > 0
-            else hl.empty_array(hl.tstr)
+            [e.s for e in row[0].entries] if len(row) > 0 else hl.empty_array(hl.tstr)
         ),
     )
     return ht.annotate(entries=ht.entries.map(lambda s: s.drop('s')))

v03_pipeline/lib/reference_data/combine.py renamed to v03_pipeline/lib/reference_data/dataset_table_operations.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from types import FunctionType
 
 import hail as hl
 import pytz
@@ -11,21 +12,75 @@
 from v03_pipeline.lib.reference_data.config import CONFIG
 
 
-def parse_version(ht: hl.Table, dataset: str, config: dict) -> hl.StringExpression:
-    annotated_version = ht.globals.get('version', hl.missing(hl.tstr))
-    config_version = config.get('version', hl.missing(hl.tstr))
-    return (
-        hl.case()
-        .when(hl.is_missing(config_version), annotated_version)
-        .when(hl.is_missing(annotated_version), config_version)
-        .when(annotated_version == config_version, config_version)
-        .or_error(
-            f'found mismatching versions for dataset {dataset}, {config_version}, {hl.eval(annotated_version)}',
+def update_or_create_joined_ht(
+    reference_dataset_collection: ReferenceDatasetCollection,
+    dataset_type: DatasetType,
+    reference_genome: ReferenceGenome,
+    dataset: str | None,
+    joined_ht: hl.Table,
+) -> hl.Table:
+    datasets = (
+        [dataset]
+        if dataset is not None
+        else reference_dataset_collection.datasets(dataset_type)
+    )
+
+    for dataset in datasets:
+        dataset_ht = get_dataset_ht(dataset, reference_genome)
+
+        if dataset in joined_ht.row:
+            joined_ht = joined_ht.drop(dataset)
+
+        joined_ht = joined_ht.join(dataset_ht, 'outer')
+        joined_ht = annotate_dataset_globals(joined_ht, dataset, dataset_ht)
+
+    return joined_ht.filter(
+        hl.any(
+            [
+                ~hl.is_missing(joined_ht[dataset])
+                for dataset in reference_dataset_collection.datasets(dataset_type)
+            ],
+        ),
+    )
+
+
+def get_dataset_ht(
+    dataset: str,
+    reference_genome: ReferenceGenome,
+) -> hl.Table:
+    config = CONFIG[dataset][reference_genome.v02_value]
+    ht = (
+        config['custom_import'](config['source_path'], reference_genome)
+        if 'custom_import' in config
+        else hl.read_table(config['path'])
+    )
+    if hasattr(ht, 'locus'):
+        ht = ht.filter(
+            hl.set(reference_genome.standard_contigs).contains(ht.locus.contig),
         )
+
+    ht = ht.filter(config['filter'](ht)) if 'filter' in config else ht
+    ht = ht.select(
+        **{
+            **get_select_fields(config.get('select'), ht),
+            **get_custom_select_fields(config.get('custom_select'), ht),
+        },
     )
+    ht = ht.transmute(**get_enum_select_fields(config.get('enum_select'), ht))
+    ht = ht.select_globals(
+        path=(config['source_path'] if 'custom_import' in config else config['path']),
+        version=parse_dataset_version(ht, dataset, config),
+        enums=hl.Struct(
+            **config.get(
+                'enum_select',
+                hl.missing(hl.tstruct(hl.tstr, hl.tarray(hl.tstr))),
+            ),
+        ),
+    )
+    return ht.select(**{dataset: ht.row.drop(*ht.key)}).distinct()
 
 
-def get_select_fields(selects, base_ht):
+def get_select_fields(selects: list | dict | None, base_ht: hl.Table) -> dict:
     """
     Generic function that takes in a select config and base_ht and generates a
     select dict that is generated from traversing the base_ht and extracting the right
@@ -57,13 +112,13 @@ def get_select_fields(selects, base_ht):
     return select_fields
 
 
-def get_custom_select_fields(custom_select, ht):
+def get_custom_select_fields(custom_select: FunctionType | None, ht: hl.Table) -> dict:
     if custom_select is None:
         return {}
     return custom_select(ht)
 
 
-def get_enum_select_fields(enum_selects, ht):
+def get_enum_select_fields(enum_selects: dict | None, ht: hl.Table) -> dict:
     enum_select_fields = {}
     if enum_selects is None:
         return enum_select_fields
@@ -89,40 +144,22 @@ def get_enum_select_fields(enum_selects, ht):
     return enum_select_fields
 
 
-def get_ht(
+def parse_dataset_version(
+    ht: hl.Table,
     dataset: str,
-    reference_genome: ReferenceGenome,
-):
-    config = CONFIG[dataset][reference_genome.v02_value]
-    ht = (
-        config['custom_import'](config['source_path'], reference_genome)
-        if 'custom_import' in config
-        else hl.read_table(config['path'])
-    )
-    if hasattr(ht, 'locus'):
-        ht = ht.filter(
-            hl.set(reference_genome.standard_contigs).contains(ht.locus.contig),
+    config: dict,
+) -> hl.StringExpression:
+    annotated_version = ht.globals.get('version', hl.missing(hl.tstr))
+    config_version = config.get('version', hl.missing(hl.tstr))
+    return (
+        hl.case()
+        .when(hl.is_missing(config_version), annotated_version)
+        .when(hl.is_missing(annotated_version), config_version)
+        .when(annotated_version == config_version, config_version)
+        .or_error(
+            f'found mismatching versions for dataset {dataset}, {config_version}, {hl.eval(annotated_version)}',
         )
-
-    ht = ht.filter(config['filter'](ht)) if 'filter' in config else ht
-    ht = ht.select(
-        **{
-            **get_select_fields(config.get('select'), ht),
-            **get_custom_select_fields(config.get('custom_select'), ht),
-        },
     )
-    ht = ht.transmute(**get_enum_select_fields(config.get('enum_select'), ht))
-    ht = ht.select_globals(
-        path=(config['source_path'] if 'custom_import' in config else config['path']),
-        version=parse_version(ht, dataset, config),
-        enums=hl.Struct(
-            **config.get(
-                'enum_select',
-                hl.missing(hl.tstruct(hl.tstr, hl.tarray(hl.tstr))),
-            ),
-        ),
-    )
-    return ht.select(**{dataset: ht.row.drop(*ht.key)}).distinct()
 
 
 def annotate_dataset_globals(joined_ht: hl.Table, dataset: str, dataset_ht: hl.Table):
@@ -153,7 +190,7 @@ def join_hts(
         ),
     )
     for dataset in reference_dataset_collection.datasets(dataset_type):
-        dataset_ht = get_ht(dataset, reference_genome)
+        dataset_ht = get_dataset_ht(dataset, reference_genome)
         joined_ht = joined_ht.join(dataset_ht, 'outer')
         joined_ht = annotate_dataset_globals(joined_ht, dataset, dataset_ht)
     return joined_ht
@@ -167,7 +204,7 @@ def update_existing_joined_hts(
     reference_dataset_collection: ReferenceDatasetCollection,
 ):
     joined_ht = hl.read_table(destination_path)
-    dataset_ht = get_ht(dataset, reference_genome)
+    dataset_ht = get_dataset_ht(dataset, reference_genome)
     joined_ht = joined_ht.drop(dataset)
     joined_ht = joined_ht.join(dataset_ht, 'outer')
     joined_ht = joined_ht.filter(