Fix performance for reloading saved variant annotations

hail_search/queries/base.py

@@ -738,7 +738,12 @@ def _parse_intervals(self, intervals, gene_ids=None, variant_keys=None, variant_
             return intervals
 
         if variant_ids:
-            intervals = [(chrom, pos, pos+1) for chrom, pos, _, _ in variant_ids]
+            first_chrom = variant_ids[0][0]
+            if all(first_chrom == v[0] for v in variant_ids):
+                positions = [pos for _, pos, _, _ in variant_ids]
+                intervals = [(first_chrom, min(positions), max(positions) + 1)]
+            else:
+                intervals = [(chrom, pos, pos+1) for chrom, pos, _, _ in variant_ids]
 
         is_x_linked = self._inheritance_mode == X_LINKED_RECESSIVE
         if not (intervals or is_x_linked):

seqr/management/commands/check_for_new_samples_from_pipeline.py

@@ -336,14 +336,21 @@ def _reload_shared_variant_annotations(data_type, genome_version, updated_varian
             variant_id: {k: v for k, v in variant.items() if k not in {'familyGuids', 'genotypes'}}
             for variant_id, variant in (updated_variants_by_id or {}).items()
         }
-        fetch_variant_ids = sorted(set(variants_by_id.keys()) - set(updated_variants_by_id.keys()))
+        fetch_variant_ids = set(variants_by_id.keys()) - set(updated_variants_by_id.keys())
         if fetch_variant_ids:
-            if not is_sv:
-                fetch_variant_ids = [parse_valid_variant_id(variant_id) for variant_id in fetch_variant_ids]
-            for i in range(0, len(fetch_variant_ids), MAX_LOOKUP_VARIANTS):
-                updated_variants = hail_variant_multi_lookup(USER_EMAIL, fetch_variant_ids[i:i+MAX_LOOKUP_VARIANTS], data_type, genome_version)
-                logger.info(f'Fetched {len(updated_variants)} additional variants')
-                updated_variants_by_id.update({variant['variantId']: variant for variant in updated_variants})
+            if is_sv:
+                variant_ids_by_chrom = {'all': fetch_variant_ids}
+            else:
+                variant_ids_by_chrom = defaultdict(list)
+                for variant_id in fetch_variant_ids:
+                    parsed_id = parse_valid_variant_id(variant_id)
+                    variant_ids_by_chrom[parsed_id[0]].append(parsed_id)
+            for chrom, variant_ids in sorted(variant_ids_by_chrom.items()):
+                variant_ids = sorted(variant_ids)
+                for i in range(0, len(variant_ids), MAX_LOOKUP_VARIANTS):
+                    updated_variants = hail_variant_multi_lookup(USER_EMAIL, variant_ids[i:i+MAX_LOOKUP_VARIANTS], data_type, genome_version)
+                    logger.info(f'Fetched {len(updated_variants)} additional variants in chromosome {chrom}')
+                    updated_variants_by_id.update({variant['variantId']: variant for variant in updated_variants})
 
         updated_variant_models = []
         for variant_id, variant in updated_variants_by_id.items():

seqr/management/tests/check_for_new_samples_from_pipeline_tests.py

@@ -423,7 +423,7 @@ def test_command(self, mock_email, mock_airtable_utils, mock_open_write_file, mo
                 {'individual_guid': 'I000017_na20889', 'family_guid': 'F000012_12', 'project_guid': 'R0003_test', 'affected': 'A', 'sample_id': 'NA20889', 'sample_type': 'WES'},
             ]}},
         ], reload_annotations_logs=[
-            'Reloading shared annotations for 3 SNV_INDEL GRCh38 saved variants (3 unique)', 'Fetched 1 additional variants', 'Fetched 1 additional variants', 'Updated 2 SNV_INDEL GRCh38 saved variants',
+            'Reloading shared annotations for 3 SNV_INDEL GRCh38 saved variants (3 unique)', 'Fetched 1 additional variants in chromosome 1', 'Fetched 1 additional variants in chromosome 1', 'Updated 2 SNV_INDEL GRCh38 saved variants',
             'No additional SV_WES GRCh38 saved variants to update',
         ], run_loading_logs={
             'GRCh38/SNV_INDEL': 'Loading 4 WES SNV_INDEL samples in 2 projects',

seqr/management/tests/reload_saved_variant_annotations_tests.py

@@ -39,18 +39,24 @@ def test_command(self, mock_logger):
 
         mock_logger.info.assert_has_calls([mock.call(log) for log in [
             'Reloading shared annotations for 3 SNV_INDEL GRCh37 saved variants (3 unique)',
-            'Fetched 2 additional variants',
+            'Fetched 2 additional variants in chromosome 1',
+            'Fetched 2 additional variants in chromosome 21',
             'Updated 2 SNV_INDEL GRCh37 saved variants',
         ]])
 
-        self.assertEqual(len(responses.calls), 1)
-        multi_lookup_request = responses.calls[0].request
-        self.assertEqual(multi_lookup_request.url, f'{MOCK_HAIL_ORIGIN}:5000/multi_lookup')
-        self.assertEqual(multi_lookup_request.headers.get('From'), 'manage_command')
-        self.assertDictEqual(json.loads(multi_lookup_request.body), {
+        self.assertEqual(len(responses.calls), 2)
+        for call in responses.calls:
+            self.assertEqual(call.request.url, f'{MOCK_HAIL_ORIGIN}:5000/multi_lookup')
+            self.assertEqual(call.request.headers.get('From'), 'manage_command')
+        self.assertDictEqual(json.loads(responses.calls[0].request.body), {
             'genome_version': 'GRCh37',
             'data_type': 'SNV_INDEL',
-            'variant_ids': [['1', 248367227, 'TC', 'T'], ['1', 46859832, 'G', 'A'], ['21', 3343353, 'GAGA', 'G']],
+            'variant_ids': [['1', 46859832, 'G', 'A'], ['1', 248367227, 'TC', 'T']],
+        })
+        self.assertDictEqual(json.loads(responses.calls[1].request.body), {
+            'genome_version': 'GRCh37',
+            'data_type': 'SNV_INDEL',
+            'variant_ids': [['21', 3343353, 'GAGA', 'G']],
         })
 
         annotation_updated_json_1 = SavedVariant.objects.get(guid='SV0000002_1248367227_r0390_100').saved_variant_json
@@ -66,11 +72,12 @@ def test_command(self, mock_logger):
         self.assertEqual(len(annotation_updated_json_2['genotypes']), 3)
 
         # Test SVs
+        responses.calls.reset()
         Sample.objects.filter(guid='S000147_na21234').update(individual_id=20)
         call_command('reload_saved_variant_annotations', 'SV_WGS', 'GRCh37')
 
-        self.assertEqual(len(responses.calls), 2)
-        self.assertDictEqual(json.loads(responses.calls[1].request.body), {
+        self.assertEqual(len(responses.calls), 1)
+        self.assertDictEqual(json.loads(responses.calls[0].request.body), {
             'genome_version': 'GRCh37',
             'data_type': 'SV_WGS',
             'variant_ids': ['prefix_19107_DEL'],