Join new entry variants in clickhouse

loading_pipeline/lib/misc/clickhouse.py

@@ -655,9 +655,17 @@ def insert_new_entries(
         )
     ]
     common, overrides = [c for c in dst_cols if c in src_cols], {}
-    if 'geneId_ids' in dst_cols and 'geneIds' in src_cols:
-        common = [c for c in common if c not in ('geneId_ids', 'geneIds')]
+    if 'xpos' not in common:
+        common.append('xpos')
+        overrides['xpos'] = 'v.xpos'
+
+    if 'geneId_ids' in dst_cols:
         common.append('geneId_ids')
+        gene_list_field = (
+            'sortedGeneConsequences'
+            if table_name_builder.dataset_type == DatasetType.SV
+            else 'sortedTranscriptConsequences'
+        )
         overrides['geneId_ids'] = f"""
             arrayFilter(
                 x -> x IS NOT NULL,
@@ -667,7 +675,7 @@ def insert_new_entries(
                         'seqrdb_id',
                         g
                     ),
-                    geneIds
+                    arrayDistinct(v.{gene_list_field}.geneId)
                 )
             )
         """
@@ -682,12 +690,21 @@ def insert_new_entries(
         """
 
     dst_list = ', '.join(common)
-    src_list = ', '.join([overrides.get(c, c) for c in common])
+    src_list = ', '.join([overrides.get(c, f'e.{c}') for c in common])
     logged_query(
         f"""
         INSERT INTO {table_name_builder.staging_dst_table(ClickHouseTable.ENTRIES)} ({dst_list})
-        SELECT {src_list}
-        FROM {table_name_builder.src_table(ClickHouseTable.ENTRIES)}
+        SELECT e.key, {src_list}
+        FROM (
+            SELECT
+                dst.key,
+                COLUMNS('.*') EXCEPT(variantId, key)
+            FROM {table_name_builder.src_table(ClickHouseTable.ENTRIES)} src
+            INNER JOIN {table_name_builder.dst_table(ClickHouseTable.KEY_LOOKUP)} dst
+            ON {ClickHouseTable.KEY_LOOKUP.join_condition}
+        ) e
+        INNER JOIN {table_name_builder.dst_table(ClickHouseTable.VARIANTS_MEMORY)} v
+        ON assumeNotNull(e.key) = v.key
         """,  # nosec B608
     )
 
@@ -952,28 +969,28 @@ def export_existing_variants_to_parquet(
     reference_genome: ReferenceGenome,
     dataset_type: DatasetType,
     run_id: str,
-    export_select_fields: str,
 ) -> None:
     table_name_builder = TableNameBuilder(
         reference_genome,
         dataset_type,
         run_id,
     )
     variants_table = table_name_builder.dst_table(
-        ClickHouseTable.VARIANT_DETAILS
+        ClickHouseTable.KEY_LOOKUP
         if dataset_type.should_write_new_variant_details
         else ClickHouseTable.VARIANTS_MEMORY,
     )
     export_table = table_name_builder.src_table(
         ClickHouseTable.EXISTING_VARIANTS,
     ).replace(
-        '/*.parquet',
+        '/*.parquet.gz',
         '',
     )
+    dt_fields = ', end, endChrom' if dataset_type == DatasetType.SV else ''
     logged_query(
         f"""
         INSERT INTO FUNCTION {export_table}
-        SELECT {export_select_fields}
+        SELECT key AS key_, variantId AS variant_id {dt_fields}
         FROM {variants_table}
         """,  # nosec B608
     )

loading_pipeline/lib/tasks/exports/fields.py

@@ -1,5 +1,6 @@
 import hail as hl
 
+from loading_pipeline.lib.annotations.shared import variant_id, xpos
 from loading_pipeline.lib.core import DatasetType, ReferenceGenome, SampleType
 from loading_pipeline.lib.tasks.exports.misc import (
     reformat_transcripts_for_export,
@@ -90,17 +91,7 @@ def get_dataset_type_specific_variants_annotations(
     }[dataset_type](ht)
 
 
-def get_existing_variants_export_field(dataset_type: DatasetType) -> str:
-    dt_fields = {
-        DatasetType.SNV_INDEL: ', transcripts.geneId AS geneIds',
-        DatasetType.MITO: '',
-        DatasetType.SV: ', CAST(xpos AS Int64) AS xpos, end, endChrom, sortedGeneConsequences.geneId AS geneIds',
-        DatasetType.GCNV: ', CAST(xpos AS Int64) AS xpos, pos AS start, end, numExon as num_exon, sortedGeneConsequences.geneId AS gene_ids',
-    }[dataset_type]
-    return f'key AS key_, variantId AS variant_id {dt_fields}'
-
-
-def get_entries_call_annotations_fields(
+def _get_entries_call_annotations_fields(
     dataset_type: DatasetType,
 ):
     if dataset_type == DatasetType.GCNV:
@@ -154,7 +145,7 @@ def _get_calls_export_fields(
                     getattr(fe, f'sample_{field}'),
                     getattr(ht, field),
                 )
-                for field in get_entries_call_annotations_fields(dataset_type)
+                for field in _get_entries_call_annotations_fields(dataset_type)
             },
             newCall=fe.concordance.new_call,
             prevCall=fe.concordance.prev_call,
@@ -163,20 +154,6 @@ def _get_calls_export_fields(
     }[dataset_type](fe)
 
 
-def get_entries_annotations_export_fields(dataset_type: DatasetType):
-    fields = {
-        'key_': lambda ht: ht.key_,
-        'xpos': lambda ht: hl.int64(ht.xpos),
-    }
-    if dataset_type in {DatasetType.SV, DatasetType.SNV_INDEL}:
-        fields['geneIds'] = lambda ht: (
-            hl.set(ht.sorted_gene_consequences.gene_id)
-            if dataset_type == DatasetType.SV
-            else hl.set(ht.sorted_transcript_consequences.gene_id)
-        )
-    return fields
-
-
 def get_entries_export_fields(
     ht: hl.Table,
     dataset_type: DatasetType,
@@ -185,16 +162,14 @@ def get_entries_export_fields(
     return {
         'project_guid': ht.family_entries.project_guid[0],
         'family_guid': ht.family_entries.family_guid[0],
-        **{
-            field: getattr(ht, field)
-            for field in get_entries_annotations_export_fields(dataset_type)
-        },
         **(
             {
                 'sample_type': sample_type.value,
+                'variantId': variant_id(ht),
+                'xpos': xpos(ht),
             }
             if dataset_type in {DatasetType.SNV_INDEL, DatasetType.MITO}
-            else {}
+            else {'variantId': ht.variant_id}
         ),
         'filters': ht.filters,
         'calls': hl.sorted(ht.family_entries, key=lambda fe: fe.s).map(

loading_pipeline/lib/tasks/exports/write_new_entries_parquet.py

@@ -3,13 +3,11 @@
 import luigi.util
 
 from loading_pipeline.lib.annotations.fields import get_fields
-from loading_pipeline.lib.annotations.shared import xpos
 from loading_pipeline.lib.misc.family_entries import (
     compute_callset_family_entries_ht,
     deduplicate_by_most_non_ref_calls,
     deglobalize_ids,
 )
-from loading_pipeline.lib.misc.io import import_parquet
 from loading_pipeline.lib.paths import (
     new_entries_parquet_path,
 )
@@ -18,17 +16,9 @@
 )
 from loading_pipeline.lib.tasks.base.base_write_parquet import BaseWriteParquetTask
 from loading_pipeline.lib.tasks.exports.fields import (
-    get_entries_annotations_export_fields,
-    get_entries_call_annotations_fields,
     get_entries_export_fields,
 )
 from loading_pipeline.lib.tasks.files import GCSorLocalTarget
-from loading_pipeline.lib.tasks.write_existing_variants_parquet import (
-    WriteExistingVariantsParquetTask,
-)
-from loading_pipeline.lib.tasks.write_new_variants_table import (
-    WriteNewVariantsTableTask,
-)
 from loading_pipeline.lib.tasks.write_remapped_and_subsetted_callset import (
     WriteRemappedAndSubsettedCallsetTask,
 )
@@ -47,45 +37,11 @@ def output(self) -> luigi.Target:
 
     def requires(self) -> list[luigi.Task]:
         return [
-            self.clone(WriteExistingVariantsParquetTask),
-            self.clone(WriteNewVariantsTableTask),
             self.clone(WriteRemappedAndSubsettedCallsetTask),
         ]
 
     def create_table(self) -> hl.Table:
-        annotations_ht = hl.read_table(self.input()[1].path)
-        annotation_selects = {
-            field: func(annotations_ht)
-            for field, func in {
-                **get_entries_annotations_export_fields(self.dataset_type),
-                **get_entries_call_annotations_fields(self.dataset_type),
-            }.items()
-        }
-        annotations_ht = annotations_ht.select(**annotation_selects)
-
-        existing_annotations_ht = import_parquet(
-            self.input()[0].path,
-            self.reference_genome,
-            self.dataset_type,
-        )
-        if 'xpos' not in existing_annotations_ht.row:
-            existing_annotations_ht = existing_annotations_ht.annotate(
-                xpos=hl.int64(xpos(existing_annotations_ht)),
-            )
-        if 'gene_ids' in existing_annotations_ht.row:
-            existing_annotations_ht = existing_annotations_ht.annotate(
-                gene_ids=hl.set(existing_annotations_ht.gene_ids),
-            )
-        if 'geneIds' in existing_annotations_ht.row:
-            existing_annotations_ht = existing_annotations_ht.annotate(
-                geneIds=hl.set(existing_annotations_ht.geneIds),
-            )
-
-        annotations_ht = annotations_ht.union(
-            existing_annotations_ht.select(*annotation_selects),
-        )
-
-        mt = hl.read_matrix_table(self.input()[2].path)
+        mt = hl.read_matrix_table(self.input()[0].path)
         ht = compute_callset_family_entries_ht(
             self.dataset_type,
             mt,
@@ -97,7 +53,6 @@ def create_table(self) -> hl.Table:
         )
         ht = deglobalize_ids(ht)
         ht = deduplicate_by_most_non_ref_calls(ht)
-        ht = ht.join(annotations_ht)
 
         # the family entries ht will contain rows
         # where at least one family is defined... after explosion,

loading_pipeline/lib/tasks/write_existing_variants_parquet.py

@@ -6,7 +6,6 @@
 from loading_pipeline.lib.tasks.base.base_loading_run_params import (
     BaseLoadingRunParams,
 )
-from loading_pipeline.lib.tasks.exports.fields import get_existing_variants_export_field
 from loading_pipeline.lib.tasks.files import (
     GCSorLocalTarget,
 )
@@ -27,10 +26,8 @@ def complete(self) -> bool:
         return self.output().exists()
 
     def run(self):
-        export_select_fields = get_existing_variants_export_field(self.dataset_type)
         export_existing_variants_to_parquet(
             self.reference_genome,
             self.dataset_type,
             self.run_id,
-            export_select_fields,
         )