diff --git a/loading_pipeline/lib/misc/clickhouse.py b/loading_pipeline/lib/misc/clickhouse.py index 480356abbc..3a5189a876 100644 --- a/loading_pipeline/lib/misc/clickhouse.py +++ b/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,7 +969,6 @@ 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, @@ -960,20 +976,21 @@ def export_existing_variants_to_parquet( 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 ) diff --git a/loading_pipeline/lib/tasks/exports/fields.py b/loading_pipeline/lib/tasks/exports/fields.py index a29ba5b95c..020b92915f 100644 --- a/loading_pipeline/lib/tasks/exports/fields.py +++ b/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( diff --git a/loading_pipeline/lib/tasks/exports/write_new_entries_parquet.py b/loading_pipeline/lib/tasks/exports/write_new_entries_parquet.py index 810c8e56f8..dfd4402803 100644 --- a/loading_pipeline/lib/tasks/exports/write_new_entries_parquet.py +++ b/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, diff --git a/loading_pipeline/lib/tasks/write_existing_variants_parquet.py b/loading_pipeline/lib/tasks/write_existing_variants_parquet.py index da89e0ce79..d108128428 100644 --- a/loading_pipeline/lib/tasks/write_existing_variants_parquet.py +++ b/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, )