diff --git a/loading_pipeline/docs/Diagrams.md b/loading_pipeline/docs/Diagrams.md index 95541d3b3a..e05ea1f382 100644 --- a/loading_pipeline/docs/Diagrams.md +++ b/loading_pipeline/docs/Diagrams.md @@ -26,32 +26,33 @@ WriteRemappedAndSubsettedCallsetTask | v - WriteMetadataForRunTask WriteExistingVariantsParquetTask - |___________________________________________| + WriteMetadataForRunTask | - v - WriteNewVariantsTableTask - | - ______________________+_______________________ - | | | - v v v - WriteNewEntries... WriteNewVariants... WriteNewVariantDetails... - ParquetTask ParquetTask ParquetTask - | | (optional) - |___________________________|_______________________| - | - v - RunPipelineTask - (all parquets ready) - | - v - WriteSuccessFileTask - | - v - LoadClickhouseVariants - | - v - LoadClickhouseEntries + _________________+__________ WriteExistingVariantsParquetTask + | | | + | |______________________| + | | + v v + WriteNewEntriesParquetTask WriteNewVariantsTableTask + | | + | ____________+____________ + | | | + | v v + | WriteNewVariantsParquetTask WriteNewVariantDetailsParquetTask + | | (optional) + |__________________________|_______________________| + | + v + RunPipelineTask (all parquets ready) + | + v + WriteSuccessFileTask + | + v + LoadClickhouseVariants + | + v + LoadClickhouseEntries ``` ClickHouse LSM-Tree diff --git a/loading_pipeline/lib/annotations/gcnv.py b/loading_pipeline/lib/annotations/gcnv.py index d9e2a8a603..76696fd54c 100644 --- a/loading_pipeline/lib/annotations/gcnv.py +++ b/loading_pipeline/lib/annotations/gcnv.py @@ -11,10 +11,6 @@ from loading_pipeline.lib.misc.gcnv import parse_gcnv_genes -def _start_and_end_equal(mt: hl.MatrixTable) -> hl.BooleanExpression: - return (mt.sample_start == mt.start) & (mt.sample_end == mt.end) - - def CN(mt: hl.MatrixTable, **_: Any) -> hl.Expression: # noqa: N802 return mt.CN @@ -104,29 +100,19 @@ def rg37_locus_end( def sample_end(mt: hl.MatrixTable, **_: Any) -> hl.Expression: - return hl.or_missing( - ~_start_and_end_equal(mt), - mt.sample_end, - ) + return mt.sample_end def sample_gene_ids(mt: hl.MatrixTable, **_: Any) -> hl.Expression: - parsed_genes = parse_gcnv_genes(mt.genes_any_overlap_Ensemble_ID) - return hl.or_missing(parsed_genes != mt.gene_ids, parsed_genes) + return parse_gcnv_genes(mt.genes_any_overlap_Ensemble_ID) def sample_start(mt: hl.MatrixTable, **_: Any) -> hl.Expression: - return hl.or_missing( - ~_start_and_end_equal(mt), - mt.sample_start, - ) + return mt.sample_start def sample_num_exon(mt: hl.MatrixTable, **_: Any) -> hl.Expression: - return hl.or_missing( - mt.genes_any_overlap_totalExons != mt.num_exon, - mt.genes_any_overlap_totalExons, - ) + return mt.genes_any_overlap_totalExons def sorted_gene_consequences( diff --git a/loading_pipeline/lib/misc/clickhouse.py b/loading_pipeline/lib/misc/clickhouse.py index 480356abbc..f35df5d3f5 100644 --- a/loading_pipeline/lib/misc/clickhouse.py +++ b/loading_pipeline/lib/misc/clickhouse.py @@ -655,9 +655,18 @@ 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')] + common.insert(0, 'key') + 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 +676,7 @@ def insert_new_entries( 'seqrdb_id', g ), - geneIds + arrayDistinct(v.{gene_list_field}.geneId) ) ) """ @@ -678,16 +687,25 @@ def insert_new_entries( ): common.append('is_gnomad_gt_5_percent') overrides['is_gnomad_gt_5_percent'] = f""" - dictGetOrDefault({ClickhouseReferenceDataset.GNOMAD_GENOMES.search_path(table_name_builder)}, 'filter_af', key, 0) > 0.05 + dictGetOrDefault({ClickhouseReferenceDataset.GNOMAD_GENOMES.search_path(table_name_builder)}, 'filter_af', e.key, 0) > 0.05 """ 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)} + 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 +970,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,7 +977,7 @@ 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, ) @@ -970,11 +987,13 @@ def export_existing_variants_to_parquet( '/*.parquet', '', ) + 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} + SETTINGS output_format_parquet_use_custom_encoder=1 """, # nosec B608 ) diff --git a/loading_pipeline/lib/misc/clickhouse_test.py b/loading_pipeline/lib/misc/clickhouse_test.py index 1126671b17..bd2cb9a2ee 100644 --- a/loading_pipeline/lib/misc/clickhouse_test.py +++ b/loading_pipeline/lib/misc/clickhouse_test.py @@ -203,7 +203,11 @@ def setUpClass(cls): with connections['clickhouse_write'].cursor() as cursor: cursor.execute( f'INSERT INTO {Env.CLICKHOUSE_DATABASE}.`seqrdb_gene_ids_src` VALUES', - [('GENE1', 123), ('GENE2', 12), ('GENE3', 1)], + [ + ('ENSG00000141510', 123), + ('ENSG00000012048', 12), + ('ENSG00000139618', 1), + ], ) cursor.execute( f'SYSTEM RELOAD DICTIONARY {Env.CLICKHOUSE_DATABASE}.`seqrdb_gene_ids`', @@ -211,6 +215,10 @@ def setUpClass(cls): def setUp(self): super().setUp() + sleep_patch = patch('time.sleep') + sleep_patch.start() + self.addCleanup(sleep_patch.stop) + base_path = runs_path( ReferenceGenome.GRCh38, DatasetType.SNV_INDEL, @@ -273,7 +281,7 @@ def write_test_parquet(df: pd.DataFrame, parquet_path: str, schema=None): # New Variants parquet. df = pd.DataFrame( { - 'key': [10, 11, 12, 13], + 'key': [20, 11, 12, 13], 'variantId': [ '1-3-A-C', '2-4-A-T', @@ -290,8 +298,19 @@ def write_test_parquet(df: pd.DataFrame, parquet_path: str, schema=None): TEST_RUN_ID, ), ) + gcnv_variants_df = pd.DataFrame( + { + 'key': [10, 11, 12, 13], + 'variantId': [ + 'suffix_1000_DEL', + 'suffix_1001_DUP', + 'suffix_1002_DEL', + 'suffix_1003_DUP', + ], + }, + ) write_test_parquet( - df, + gcnv_variants_df, new_variants_parquet_path( ReferenceGenome.GRCh38, DatasetType.GCNV, @@ -302,7 +321,6 @@ def write_test_parquet(df: pd.DataFrame, parquet_path: str, schema=None): # New Entries Parquet df = pd.DataFrame( { - 'key': [10, 3, 4], 'project_guid': [ 'project_d', 'project_d', @@ -323,10 +341,10 @@ def write_test_parquet(df: pd.DataFrame, parquet_path: str, schema=None): 'WES', 'WES', ], - 'geneIds': [ - [], - ['GENE1', 'GENE2'], - ['GENE3'], + 'variantId': [ + '10-987654-G-A', + '3-133456789-A-G', + '4-133456789-C-T', ], 'calls': [ [('sample_d1', 0), ('sample_d11', 2)], @@ -342,12 +360,11 @@ def write_test_parquet(df: pd.DataFrame, parquet_path: str, schema=None): ) schema = pa.schema( [ - ('key', pa.int64()), ('project_guid', pa.string()), ('family_guid', pa.string()), ('xpos', pa.int64()), ('sample_type', pa.string()), - ('geneIds', pa.list_(pa.string())), + ('variantId', pa.string()), ( 'calls', pa.list_( @@ -366,14 +383,101 @@ def write_test_parquet(df: pd.DataFrame, parquet_path: str, schema=None): ), schema, ) + gcnv_entries_df = pd.DataFrame( + { + 'project_guid': [ + 'project_d', + 'project_d', + 'project_d', + ], + 'family_guid': [ + 'family_d1', + 'family_d2', + 'family_d3', + ], + 'variantId': [ + 'suffix_1000_DEL', + 'suffix_1002_DEL', + 'suffix_1003_DUP', + ], + 'calls': [ + [ + { + 'sampleId': 'sample_d1', + 'gt': 0, + 'cn': 2, + 'qs': 4, + 'defragged': False, + 'start': 100006937, + 'end': 100007881, + 'numExon': 2, + 'geneIds': ['ENSG00000117620', 'ENSG00000283761'], + 'newCall': False, + 'prevCall': True, + 'prevOverlap': False, + }, + { + 'sampleId': 'sample_d11', + 'gt': 2, + 'cn': 0, + 'qs': 30, + 'defragged': False, + 'start': 100006937, + 'end': 100007881, + 'numExon': 2, + 'geneIds': ['ENSG00000117620', 'ENSG00000283761'], + 'newCall': True, + 'prevCall': False, + 'prevOverlap': False, + }, + ], + [ + { + 'sampleId': 'sample_d2', + 'gt': 0, + 'cn': 2, + 'qs': 5, + 'defragged': False, + 'start': 100017585, + 'end': 100023213, + 'numExon': 1, + 'geneIds': ['ENSG00000117620', 'ENSG00000283761'], + 'newCall': False, + 'prevCall': True, + 'prevOverlap': False, + }, + ], + [ + { + 'sampleId': 'sample_d3', + 'gt': 1, + 'cn': 1, + 'qs': 20, + 'defragged': False, + 'start': 100017585, + 'end': 100023213, + 'numExon': 1, + 'geneIds': ['ENSG00000117620', 'ENSG00000283761'], + 'newCall': True, + 'prevCall': False, + 'prevOverlap': False, + }, + ], + ], + 'sign': [ + 1, + 1, + 1, + ], + }, + ) write_test_parquet( - df.drop('geneIds', axis=1), + gcnv_entries_df, new_entries_parquet_path( ReferenceGenome.GRCh38, DatasetType.GCNV, TEST_RUN_ID, ), - schema.remove(5).remove(5), ) def test_get_clickhouse_client(self): @@ -973,6 +1077,7 @@ def test_load_run_variants_snv_indel(self, mock_for_reference_genome_dataset_typ SELECT * FROM {Env.CLICKHOUSE_DATABASE}.`GRCh38/SNV_INDEL/variants_memory` + WHERE key > 5 """, # nosec B608 ) variants_memory = cursor.fetchall() @@ -983,6 +1088,7 @@ def test_load_run_variants_snv_indel(self, mock_for_reference_genome_dataset_typ (11, [], [], []), (12, [], [], []), (13, [], [], []), + (20, [], [], []), ], ) cursor.execute( @@ -990,13 +1096,14 @@ def test_load_run_variants_snv_indel(self, mock_for_reference_genome_dataset_typ SELECT * FROM {Env.CLICKHOUSE_DATABASE}.`GRCh38/SNV_INDEL/variants_disk` + WHERE key > 5 """, # nosec B608 ) variants_disk = cursor.fetchall() self.assertCountEqual( variants_disk, [ - (10, [], [], []), + (20, [], [], []), (11, [], [], []), (12, [], [], []), (13, [], [], []), @@ -1058,6 +1165,12 @@ def test_load_run_variants_gcnv(self): self.assertEqual(key_lookup_count, 4) def test_load_run_entries_gcnv(self): + load_run_variants( + ReferenceGenome.GRCh38, + DatasetType.GCNV, + TEST_RUN_ID, + ) + load_run_entries( ReferenceGenome.GRCh38, DatasetType.GCNV, diff --git a/loading_pipeline/lib/tasks/exports/fields.py b/loading_pipeline/lib/tasks/exports/fields.py index a29ba5b95c..7e27ae9d70 100644 --- a/loading_pipeline/lib/tasks/exports/fields.py +++ b/loading_pipeline/lib/tasks/exports/fields.py @@ -1,9 +1,9 @@ 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, - snake_to_camelcase, ) STANDARD_CONTIGS = hl.set( @@ -90,31 +90,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( - dataset_type: DatasetType, -): - if dataset_type == DatasetType.GCNV: - return { - 'start': lambda ht: hl.int64(ht.start_locus.position), - 'end': lambda ht: hl.int64(ht.end_locus.position), - 'num_exon': lambda ht: ht.num_exon, - 'gene_ids': lambda ht: hl.set(ht.sorted_gene_consequences.gene_id), - } - return {} - - def _get_calls_export_fields( - ht: hl.Table, fe: hl.Struct, dataset_type: DatasetType, ): @@ -149,13 +125,10 @@ def _get_calls_export_fields( cn=fe.CN, qs=fe.QS, defragged=fe.defragged, - **{ - snake_to_camelcase(field): hl.or_else( - getattr(fe, f'sample_{field}'), - getattr(ht, field), - ) - for field in get_entries_call_annotations_fields(dataset_type) - }, + start=fe.sample_start, + end=fe.sample_end, + numExon=fe.sample_num_exon, + geneIds=fe.sample_gene_ids, newCall=fe.concordance.new_call, prevCall=fe.concordance.prev_call, prevOverlap=fe.concordance.prev_overlap, @@ -163,20 +136,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,20 +144,18 @@ 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( - lambda fe: _get_calls_export_fields(ht, fe, dataset_type), + lambda fe: _get_calls_export_fields(fe, dataset_type), ), 'sign': 1, } 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..35bb6cad41 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,16 +16,11 @@ ) 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_metadata_for_run import ( + WriteMetadataForRunTask, ) from loading_pipeline.lib.tasks.write_remapped_and_subsetted_callset import ( WriteRemappedAndSubsettedCallsetTask, @@ -47,45 +40,12 @@ def output(self) -> luigi.Target: def requires(self) -> list[luigi.Task]: return [ - self.clone(WriteExistingVariantsParquetTask), - self.clone(WriteNewVariantsTableTask), self.clone(WriteRemappedAndSubsettedCallsetTask), + self.clone(WriteMetadataForRunTask), ] 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 +57,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/exports/write_new_entries_parquet_test.py b/loading_pipeline/lib/tasks/exports/write_new_entries_parquet_test.py index 05b73e2086..b8908c1f8f 100644 --- a/loading_pipeline/lib/tasks/exports/write_new_entries_parquet_test.py +++ b/loading_pipeline/lib/tasks/exports/write_new_entries_parquet_test.py @@ -1,6 +1,5 @@ -from typing import ClassVar +import os -import hail as hl import luigi.worker import pandas as pd @@ -11,15 +10,12 @@ ) from loading_pipeline.lib.misc.validation import ALL_VALIDATIONS from loading_pipeline.lib.paths import ( + existing_variants_parquet_path, new_entries_parquet_path, - new_variants_table_path, ) from loading_pipeline.lib.tasks.exports.write_new_entries_parquet import ( WriteNewEntriesParquetTask, ) -from loading_pipeline.lib.test.clickhouse_schema_testcase import ( - ClickhouseSchemaTestCase, -) from loading_pipeline.lib.test.misc import ( convert_ndarray_to_list, copy_project_pedigree_to_mocked_dir, @@ -36,61 +32,11 @@ TEST_MITO_CALLSET = 'loading_pipeline/var/test/callsets/mito_1.mt' TEST_SV_VCF_2 = 'loading_pipeline/var/test/callsets/sv_2.vcf' TEST_GCNV_BED_FILE = 'loading_pipeline/var/test/callsets/gcnv_1.tsv' -TEST_SNV_INDEL_ANNOTATIONS = ( - 'loading_pipeline/var/test/exports/GRCh38/SNV_INDEL/annotations.ht' -) -TEST_MITO_ANNOTATIONS = 'loading_pipeline/var/test/exports/GRCh38/MITO/annotations.ht' -TEST_SV_ANNOTATIONS = 'loading_pipeline/var/test/exports/GRCh38/SV/annotations.ht' -TEST_GCNV_ANNOTATIONS = 'loading_pipeline/var/test/exports/GRCh38/GCNV/annotations.ht' TEST_RUN_ID = 'manual__2024-04-03' -class WriteNewEntriesParquetTest(MockedDatarootTestCase, ClickhouseSchemaTestCase): - fixtures: ClassVar = ['clickhouse_test'] - - def setUp(self) -> None: - super().setUp() - ht = hl.read_table( - TEST_SNV_INDEL_ANNOTATIONS, - ) - ht = ht.filter(ht.variant_id != '1-878314-G-C') - ht.write( - new_variants_table_path( - ReferenceGenome.GRCh38, - DatasetType.SNV_INDEL, - TEST_RUN_ID, - ), - ) - ht = hl.read_table( - TEST_MITO_ANNOTATIONS, - ) - ht.write( - new_variants_table_path( - ReferenceGenome.GRCh38, - DatasetType.MITO, - TEST_RUN_ID, - ), - ) - ht = hl.read_table( - TEST_SV_ANNOTATIONS, - ) - ht.write( - new_variants_table_path( - ReferenceGenome.GRCh38, - DatasetType.SV, - TEST_RUN_ID, - ), - ) - ht = hl.read_table(TEST_GCNV_ANNOTATIONS) - ht.write( - new_variants_table_path( - ReferenceGenome.GRCh38, - DatasetType.GCNV, - TEST_RUN_ID, - ), - ) - +class WriteNewEntriesParquetTest(MockedDatarootTestCase): def test_write_new_entries_parquet(self): copy_project_pedigree_to_mocked_dir( TEST_PEDIGREE_3_REMAP, @@ -128,36 +74,35 @@ def test_write_new_entries_parquet(self): ), ) export_json = convert_ndarray_to_list(df.to_dict('records')) - self.assertEqual(len(export_json), 16) + self.assertEqual(len(export_json), 181) self.assertEqual( df['family_guid'].value_counts().to_dict(), { - 'abc_1': 2, - '345_1': 2, - '123_1': 1, - '234_1': 1, - '456_1': 1, - '567_1': 1, - '678_1': 1, - '789_1': 1, - '890_1': 1, - '901_1': 1, - 'bcd_1': 1, - 'cde_1': 1, - 'def_1': 1, - 'efg_1': 1, + 'abc_1': 16, + '789_1': 15, + '890_1': 14, + '901_1': 14, + 'bcd_1': 14, + '345_1': 13, + '456_1': 13, + '567_1': 13, + 'def_1': 13, + '123_1': 12, + '234_1': 11, + '678_1': 11, + 'cde_1': 11, + 'efg_1': 11, }, ) self.assertEqual( - [export_json[0], export_json[9], export_json[15]], + [export_json[2], export_json[11], export_json[17]], [ { - 'key': 0, 'project_guid': 'R0114_project4', 'family_guid': '123_1', 'sample_type': 'WGS', + 'variantId': '1-876499-A-G', 'xpos': 1000876499, - 'geneIds': ['ENSG00000187634'], 'filters': [], 'calls': [ { @@ -171,12 +116,11 @@ def test_write_new_entries_parquet(self): 'sign': 1, }, { - 'key': 0, 'project_guid': 'R0113_test_project', 'family_guid': 'abc_1', 'sample_type': 'WGS', + 'variantId': '1-876499-A-G', 'xpos': 1000876499, - 'geneIds': ['ENSG00000187634'], 'filters': [], 'calls': [ { @@ -204,12 +148,11 @@ def test_write_new_entries_parquet(self): 'sign': 1, }, { - 'key': 1, 'project_guid': 'R0113_test_project', 'family_guid': 'abc_1', 'sample_type': 'WGS', + 'variantId': '1-878314-G-C', 'xpos': 1000878314, - 'geneIds': ['ENSG00000177000'], 'filters': ['VQSRTrancheSNP99.00to99.90'], 'calls': [ { @@ -263,14 +206,15 @@ def test_mito_write_new_entries_parquet(self): ), ) export_json = convert_ndarray_to_list(df.to_dict('records')) + self.assertEqual(len(export_json), 3) self.assertEqual( - export_json, + export_json[:1], [ { - 'key': 998, 'project_guid': 'R0116_test_project3', 'family_guid': 'family_1', 'sample_type': 'WGS', + 'variantId': 'M-8-G-T', 'xpos': 25000000008, 'filters': [], 'calls': [ @@ -296,6 +240,20 @@ def test_sv_write_new_entries_parquet(self): SampleType.WGS, 'R0115_test_project2', ) + existing_variants_path = existing_variants_parquet_path( + ReferenceGenome.GRCh38, + DatasetType.SV, + TEST_RUN_ID, + ) + os.makedirs(os.path.dirname(existing_variants_path), exist_ok=True) + pd.DataFrame( + { + 'variant_id': ['BND_chr1_6'], + 'key_': [727], + 'end': [180928], + 'endChrom': ['chr5'], + }, + ).to_parquet(existing_variants_path) worker = luigi.worker.Worker() task = WriteNewEntriesParquetTask( reference_genome=ReferenceGenome.GRCh38, @@ -318,15 +276,14 @@ def test_sv_write_new_entries_parquet(self): ), ) export_json = convert_ndarray_to_list(df.to_dict('records')) + self.assertEqual(len(export_json), 2) self.assertEqual( - export_json, + export_json[:1], [ { - 'key': 727, + 'variantId': 'BND_chr1_6', 'project_guid': 'R0115_test_project2', 'family_guid': 'family_2_1', - 'xpos': 1001025886, - 'geneIds': ['ENSG00000188157'], 'filters': ['HIGH_SR_BACKGROUND', 'UNRESOLVED'], 'calls': [ { @@ -401,14 +358,14 @@ def test_gcnv_write_new_entries_parquet(self): ), ) export_json = convert_ndarray_to_list(df.to_dict('records')) + self.assertEqual(len(export_json), 2) self.assertEqual( - export_json, + export_json[:1], [ { - 'key': 0, + 'variantId': 'suffix_16456_DEL', 'project_guid': 'R0115_test_project2', 'family_guid': 'family_2_1', - 'xpos': 1000939203, 'filters': [], 'calls': [ { @@ -433,7 +390,7 @@ def test_gcnv_write_new_entries_parquet(self): 'defragged': False, 'start': 100017585, 'end': 100023213, - 'numExon': 1, + 'numExon': 3, 'geneIds': ['ENSG00000117620', 'ENSG00000283761'], 'newCall': False, 'prevCall': False, @@ -447,7 +404,7 @@ def test_gcnv_write_new_entries_parquet(self): 'defragged': False, 'start': 100017585, 'end': 100023213, - 'numExon': 1, + 'numExon': 3, 'geneIds': ['ENSG00000117620', 'ENSG00000283761'], 'newCall': False, 'prevCall': True, diff --git a/loading_pipeline/lib/tasks/write_existing_variants_parquet.py b/loading_pipeline/lib/tasks/write_existing_variants_parquet.py index 9e9c16ee35..61adf3b70e 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, ) @@ -29,10 +28,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, ) diff --git a/loading_pipeline/lib/tasks/write_existing_variants_parquet_test.py b/loading_pipeline/lib/tasks/write_existing_variants_parquet_test.py index 1323c0b775..d3d8af7dd8 100644 --- a/loading_pipeline/lib/tasks/write_existing_variants_parquet_test.py +++ b/loading_pipeline/lib/tasks/write_existing_variants_parquet_test.py @@ -46,20 +46,18 @@ def _run_task( def test_snv_indel(self): df = self._run_task(DatasetType.SNV_INDEL) - self.assertEqual(list(df.columns), ['key_', 'variant_id', 'geneIds']) + self.assertEqual(list(df.columns), ['key_', 'variant_id']) df = df.sort_values('key_').reset_index(drop=True) self.assertEqual( convert_ndarray_to_list( - df[['key_', 'variant_id', 'geneIds']].to_dict('records'), + df[['key_', 'variant_id']].to_dict('records'), ), [ - { - 'key_': 1, - 'variant_id': '1-878314-G-C', - 'geneIds': ['ENSG00000177000'], - }, - {'key_': 7, 'variant_id': '7-1234567-AGT-A', 'geneIds': []}, - {'key_': 10, 'variant_id': '10-987654-G-A', 'geneIds': []}, + {'key_': 1, 'variant_id': '1-878314-G-C'}, + {'key_': 3, 'variant_id': '3-133456789-A-G'}, + {'key_': 4, 'variant_id': '4-133456789-C-T'}, + {'key_': 7, 'variant_id': '7-1234567-AGT-A'}, + {'key_': 10, 'variant_id': '10-987654-G-A'}, ], ) @@ -68,7 +66,7 @@ def test_grch37_snv_indel(self): DatasetType.SNV_INDEL, reference_genome=ReferenceGenome.GRCh37, ) - self.assertEqual(list(df.columns), ['key_', 'variant_id', 'geneIds']) + self.assertEqual(list(df.columns), ['key_', 'variant_id']) self.assertEqual(len(df), 0) def test_mito(self): @@ -80,7 +78,7 @@ def test_sv(self): df = self._run_task(DatasetType.SV) self.assertEqual( list(df.columns), - ['key_', 'variant_id', 'xpos', 'end', 'endChrom', 'geneIds'], + ['key_', 'variant_id', 'end', 'endChrom'], ) self.assertEqual(len(df), 0) @@ -88,6 +86,6 @@ def test_gcnv(self): df = self._run_task(DatasetType.GCNV) self.assertEqual( list(df.columns), - ['key_', 'variant_id', 'xpos', 'start', 'end', 'num_exon', 'gene_ids'], + ['key_', 'variant_id'], ) self.assertEqual(len(df), 0) diff --git a/loading_pipeline/lib/test/fixtures/clickhouse_test.json b/loading_pipeline/lib/test/fixtures/clickhouse_test.json index 9de4810176..c3d10c054e 100644 --- a/loading_pipeline/lib/test/fixtures/clickhouse_test.json +++ b/loading_pipeline/lib/test/fixtures/clickhouse_test.json @@ -302,5 +302,80 @@ "sorted_motif_feature_consequences": [], "sorted_regulatory_feature_consequences": [] } + }, + { + "model": "clickhouse_search.keylookupsnvindel", + "pk": "1-878314-G-C", + "fields": { + "variant_id": "1-878314-G-C", + "key": 1 + } + }, + { + "model": "clickhouse_search.keylookupsnvindel", + "pk": "7-1234567-AGT-A", + "fields": { + "variant_id": "7-1234567-AGT-A", + "key": 7 + } + }, + { + "model": "clickhouse_search.keylookupsnvindel", + "pk": "10-987654-G-A", + "fields": { + "variant_id": "10-987654-G-A", + "key": 10 + } + }, + { + "model": "clickhouse_search.keylookupsnvindel", + "pk": "3-133456789-A-G", + "fields": { + "variant_id": "3-133456789-A-G", + "key": 3 + } + }, + { + "model": "clickhouse_search.keylookupsnvindel", + "pk": "4-133456789-C-T", + "fields": { + "variant_id": "4-133456789-C-T", + "key": 4 + } + }, + { + "model": "clickhouse_search.variantssnvindel", + "pk": 10, + "fields": { + "key": 10, + "sorted_transcript_consequences": [], + "sorted_motif_feature_consequences": [], + "sorted_regulatory_feature_consequences": [] + } + }, + { + "model": "clickhouse_search.variantssnvindel", + "pk": 3, + "fields": { + "key": 3, + "sorted_transcript_consequences": [ + [null, null, [], null, null, "ENSG00000141510", null], + [null, null, [], null, null, "ENSG00000012048", null] + ], + "sorted_motif_feature_consequences": [], + "sorted_regulatory_feature_consequences": [] + } + }, + { + "model": "clickhouse_search.variantssnvindel", + "pk": 4, + "fields": { + "key": 4, + "sorted_transcript_consequences": [ + [null, null, [], null, null, "ENSG00000139618", null] + ], + "sorted_motif_feature_consequences": [], + "sorted_regulatory_feature_consequences": [] + } } ]