Merge pull request #5042 from broadinstitute/exclude-current-search-results

Exclude current search results from next search

Author: hanars

clickhouse_search/backend/functions.py

@@ -7,6 +7,12 @@
 class Array(Func):
     function = 'array'
 
+    def _resolve_output_field(self):
+        output_field = super()._resolve_output_field()
+        if not isinstance(output_field, ArrayField):
+            output_field = ArrayField(base_field=output_field)
+        return output_field
+
 
 class ArrayConcat(Func):
     function = 'arrayConcat'

clickhouse_search/managers.py

@@ -726,9 +726,12 @@ def clinvar_model(self):
     def genome_version(self):
         return self.annotations_model.ANNOTATION_CONSTANTS['genomeVersion']
 
-    def search(self, sample_data, freqs=None, annotations=None, **kwargs):
+    def search(self, sample_data, freqs=None, annotations=None, exclude_keys=None, **kwargs):
         entries = self.filter_locus(**kwargs)
 
+        if exclude_keys:
+            entries = entries.exclude(key__in=exclude_keys)
+
         entries = self._join_annotations(entries)
 
         is_sv_class = 'cn' in self.call_fields

clickhouse_search/search.py

@@ -34,6 +34,8 @@ def get_clickhouse_variants(samples, search, user, previous_search_results, geno
     family_guid = None
     inheritance_mode = search.get('inheritance_mode')
     has_comp_het = inheritance_mode in {RECESSIVE, COMPOUND_HET}
+    exclude_keys = search.pop('exclude_keys', None) or {}
+    exclude_key_pairs = search.pop('exclude_key_pairs', None) or {}
     for dataset_type, sample_data in sample_data_by_dataset_type.items():
         logger.info(f'Loading {dataset_type} data for {len(set(sample_data["family_guids"]))} families', user)
 
@@ -44,13 +46,13 @@ def get_clickhouse_variants(samples, search, user, previous_search_results, geno
 
         dataset_results = []
         if inheritance_mode != COMPOUND_HET:
-            result_q = _get_search_results_queryset(entry_cls, annotations_cls, sample_data, skip_individual_guid=is_multi_project, **search)
+            result_q = _get_search_results_queryset(entry_cls, annotations_cls, sample_data, skip_individual_guid=is_multi_project, exclude_keys=exclude_keys.get(dataset_type), **search)
             dataset_results += _evaluate_results(result_q)
         if has_comp_het:
             comp_het_sample_data = sample_data
             if is_multi_project and dataset_type == Sample.DATASET_TYPE_VARIANT_CALLS and _is_x_chrom_only(genome_version, **search):
                 comp_het_sample_data = _no_affected_male_families(sample_data, user)
-            result_q = _get_data_type_comp_het_results_queryset(entry_cls, annotations_cls, comp_het_sample_data, skip_individual_guid=is_multi_project, **search)
+            result_q = _get_data_type_comp_het_results_queryset(entry_cls, annotations_cls, comp_het_sample_data, exclude_key_pairs=exclude_key_pairs.get(dataset_type), skip_individual_guid=is_multi_project, **search)
             dataset_results += _evaluate_results(result_q, is_comp_het=True)
 
         if is_multi_project:
@@ -60,7 +62,7 @@ def get_clickhouse_variants(samples, search, user, previous_search_results, geno
     if has_comp_het and Sample.DATASET_TYPE_VARIANT_CALLS in sample_data_by_dataset_type and any(
         dataset_type.startswith(Sample.DATASET_TYPE_SV_CALLS) for dataset_type in sample_data_by_dataset_type
     ):
-        results += _get_multi_data_type_comp_het_results_queryset(genome_version, sample_data_by_dataset_type, user, **search)
+        results += _get_multi_data_type_comp_het_results_queryset(genome_version, sample_data_by_dataset_type, user, exclude_key_pairs, **search)
 
     cache_results = get_clickhouse_cache_results(results, sort, family_guid)
     previous_search_results.update(cache_results)
@@ -83,7 +85,7 @@ def _evaluate_results(result_q, is_comp_het=False):
         raise InvalidSearchException('This search returned too many results')
     return results
 
-def _get_multi_data_type_comp_het_results_queryset(genome_version, sample_data_by_dataset_type, user, annotations=None, annotations_secondary=None, inheritance_mode=None, **search_kwargs):
+def _get_multi_data_type_comp_het_results_queryset(genome_version, sample_data_by_dataset_type, user, exclude_key_pairs, annotations=None, annotations_secondary=None, inheritance_mode=None, **search_kwargs):
     if annotations_secondary:
         annotations = {
             **annotations,
@@ -124,7 +126,7 @@ def _get_multi_data_type_comp_het_results_queryset(genome_version, sample_data_b
         sv_annotations_cls = ANNOTATIONS_CLASS_MAP[genome_version][sv_dataset_type]
         sv_q = sv_annotations_cls.objects.subquery_join(sv_entries).search(**search_kwargs, annotations=annotations)
 
-        result_q = _get_comp_het_results_queryset(annotations_cls, snv_indel_q, sv_q, len(families))
+        result_q = _get_comp_het_results_queryset(annotations_cls, snv_indel_q, sv_q, len(families), exclude_key_pairs.get(f'{Sample.DATASET_TYPE_VARIANT_CALLS},{sv_dataset_type}'))
         dataset_results = _evaluate_results(result_q, is_comp_het=True)
         if skip_individual_guid:
             _add_individual_guids(dataset_results, sv_sample_data)
@@ -133,7 +135,7 @@ def _get_multi_data_type_comp_het_results_queryset(genome_version, sample_data_b
     return results
 
 
-def _get_data_type_comp_het_results_queryset(entry_cls, annotations_cls, sample_data, annotations=None, annotations_secondary=None, inheritance_mode=None, **search_kwargs):
+def _get_data_type_comp_het_results_queryset(entry_cls, annotations_cls, sample_data, annotations=None, annotations_secondary=None, inheritance_mode=None, exclude_key_pairs=None, **search_kwargs):
     entries = entry_cls.objects.search(
         sample_data, **search_kwargs, inheritance_mode=COMPOUND_HET, annotations=annotations, annotate_carriers=True,
     )
@@ -143,10 +145,10 @@ def _get_data_type_comp_het_results_queryset(entry_cls, annotations_cls, sample_
         annotations=annotations_secondary or annotations, **search_kwargs,
     )
 
-    return _get_comp_het_results_queryset(annotations_cls, primary_q, secondary_q, len(set(sample_data['family_guids'])))
+    return _get_comp_het_results_queryset(annotations_cls, primary_q, secondary_q, len(set(sample_data['family_guids'])), exclude_key_pairs)
 
 
-def _get_comp_het_results_queryset(annotations_cls, primary_q, secondary_q, num_families):
+def _get_comp_het_results_queryset(annotations_cls, primary_q, secondary_q, num_families, exclude_key_pairs):
     results = annotations_cls.objects.search_compound_hets(primary_q, secondary_q)
 
     if results.has_annotation('primary_carriers') and results.has_annotation('secondary_carriers'):
@@ -206,9 +208,12 @@ def _get_comp_het_results_queryset(annotations_cls, primary_q, secondary_q, num_
             **genotype_expressions,
         )
 
-    return results.annotate(
+    pair_results = results.annotate(
         pair_key=ArraySort(Array('primary_key', 'secondary_key')),
-    ).distinct('pair_key').values_list(
+    ).distinct('pair_key')
+    if exclude_key_pairs:
+        pair_results = pair_results.exclude(pair_key__in=exclude_key_pairs)
+    return pair_results.values_list(
         'pair_key',
         _result_as_tuple(results, 'primary_'),
         _result_as_tuple(results, 'secondary_'),

clickhouse_search/search_tests.py

@@ -503,6 +503,34 @@ def test_inheritance_filter(self):
             [GRCH37_VARIANT], inheritance_mode=inheritance_mode, inheritance_filter={'allowNoCall': True},
         )
 
+    def test_exclude_previous_search_results(self):
+        VariantSearchResults.objects.create(variant_search_id=79516, search_hash='abc1234')
+        self.mock_redis.get.side_effect = [None, json.dumps({'all_results': [
+            VARIANT1, VARIANT2, [VARIANT3, VARIANT2], [GCNV_VARIANT4, GCNV_VARIANT3],
+        ]})]
+        self.mock_redis.keys.side_effect = [[], ['search_results__abc1234__gnomad']]
+
+        self._assert_expected_search(
+            [[MULTI_DATA_TYPE_COMP_HET_VARIANT2, GCNV_VARIANT4], [VARIANT3, VARIANT4], GCNV_VARIANT3, MITO_VARIANT3],
+            inheritance_mode='recessive', **COMP_HET_ALL_PASS_FILTERS,
+            exclude={'previousSearch': True, 'previousSearchHash': 'abc1234'}, cached_variant_fields=[
+                [{'selectedGeneId': 'ENSG00000277258'}, {'selectedGeneId': 'ENSG00000277258'}],
+                [{'selectedGeneId': 'ENSG00000097046'}, {'selectedGeneId': 'ENSG00000097046'}],
+                {}, {},
+            ],
+        )
+
+        self.mock_redis.get.side_effect = [None, json.dumps({'all_results': [
+            [MULTI_DATA_TYPE_COMP_HET_VARIANT2, GCNV_VARIANT4], [VARIANT3, VARIANT4], GCNV_VARIANT3, MITO_VARIANT3,
+        ]})]
+        self.mock_redis.keys.side_effect = [[], ['search_results__abc1234__gnomad']]
+        self._assert_expected_search(
+            [VARIANT2, [GCNV_VARIANT3, GCNV_VARIANT4]],
+            inheritance_mode='recessive', cached_variant_fields=[
+                {}, [{'selectedGeneId': 'ENSG00000275023'}, {'selectedGeneId': 'ENSG00000275023'}],
+            ],
+        )
+
     def test_quality_filter(self):
         quality_filter = {'vcf_filter': 'pass'}
         self._assert_expected_search(

seqr/utils/search/search_utils_tests.py

@@ -279,9 +279,9 @@ def test_invalid_search_query_variants(self):
 
         self._test_invalid_search_params(query_variants)
 
-    def _test_expected_search_call(self, mock_get_variants, results_cache, search_fields=None, has_gene_search=False,
+    def _get_expected_search_call(self, results_cache, search_fields=None, has_gene_search=False,
                                    rs_ids=None, variant_ids=None, parsed_variant_ids=None, inheritance_mode='de_novo',
-                                   dataset_type=None, secondary_dataset_type=None, omitted_sample_guids=None,
+                                   dataset_type=None, secondary_dataset_type=None, exclude_keys=None, exclude_key_pairs=None,
                                    exclude_locations=False, exclude=None, annotations=None, annotations_secondary=None, single_gene_search=False, **kwargs):
         genes = intervals = None
         has_included_gene_search = has_gene_search and not exclude_locations
@@ -315,8 +315,16 @@ def _test_expected_search_call(self, mock_get_variants, results_cache, search_fi
             expected_search['annotations'] = annotations
         if annotations_secondary is not None:
             expected_search['annotations_secondary'] = annotations_secondary
+        if exclude_keys is not None:
+            expected_search['exclude_keys'] = exclude_keys
+        if exclude_key_pairs is not None:
+            expected_search['exclude_key_pairs'] = exclude_key_pairs
 
-        mock_get_variants.assert_called_with(mock.ANY, expected_search, self.user, results_cache, '37', **kwargs)
+        return (mock.ANY, expected_search, self.user, results_cache, '37'), kwargs, genes
+
+    def _test_expected_search_call(self, mock_get_variants, *args, has_gene_search=False, exclude_locations=False, omitted_sample_guids=None, **kwargs):
+        call_args, call_kwargs, genes = self._get_expected_search_call(*args, has_gene_search=has_gene_search, exclude_locations=exclude_locations, **kwargs)
+        mock_get_variants.assert_called_with(*call_args, **call_kwargs)
         self._assert_expected_search_samples(mock_get_variants, omitted_sample_guids, has_gene_search and not exclude_locations)
 
         if genes:
@@ -341,16 +349,17 @@ def _assert_expected_search_samples(self, mock_get_variants, omitted_sample_guid
 
 
     def test_query_variants(self, mock_get_variants):
+        parsed_variants = [{**v, 'key': (i+1) * 1000 } for i, v in enumerate(PARSED_VARIANTS)]
         def _mock_get_variants(families, search, user, previous_search_results, genome_version, **kwargs):
-            previous_search_results['all_results'] = PARSED_VARIANTS
+            previous_search_results['all_results'] =parsed_variants
             previous_search_results['total_results'] = 5
-            return PARSED_VARIANTS
+            return parsed_variants
         mock_get_variants.side_effect = _mock_get_variants
 
         variants, total = query_variants(self.results_model, user=self.user)
-        self.assertListEqual(variants, PARSED_VARIANTS)
+        self.assertListEqual(variants, parsed_variants)
         self.assertEqual(total, 5)
-        results_cache = {'all_results': PARSED_VARIANTS, 'total_results': 5}
+        results_cache = {'all_results': parsed_variants, 'total_results': 5}
         self.assert_cached_results(results_cache)
         self._test_expected_search_call(
             mock_get_variants, results_cache, sort='xpos', page=1, num_results=100, skip_genotype_filter=False,
@@ -466,6 +475,25 @@ def _mock_get_variants(families, search, user, previous_search_results, genome_v
             omitted_sample_guids=['S000145_hg00731', 'S000146_hg00732', 'S000148_hg00733'],
         )
 
+        self.set_cache(None)
+        mock_get_variants.reset_mock()
+        self.search_model.search = {
+            'inheritance': {'mode': 'any_affected'},
+            'exclude': {'previousSearch': True, 'previousSearchHash': 'abc1234', 'clinvar': ['benign']},
+        }
+        previous_search_model = VariantSearch.objects.create(search={'inheritance': {'mode': 'de_novo'}})
+        previous_results_model = VariantSearchResults.objects.create(variant_search=previous_search_model, search_hash='abc1234')
+        previous_results_model.families.set(self.families)
+        query_variants(self.results_model, user=self.user)
+        self._test_exclude_previous_search(
+            mock_get_variants, results_cache, sort='xpos', page=1, num_results=100, skip_genotype_filter=False,
+            inheritance_mode='any_affected', exclude={'clinvar': ['benign']}, search_fields=['exclude'],
+        )
+
+    def _test_exclude_previous_search(self, mock_get_variants, *args, num_searches=1, **kwargs):
+        self._test_expected_search_call(mock_get_variants, *args, **kwargs)
+        self.assertEqual(mock_get_variants.call_count, num_searches)
+
     def test_cached_query_variants(self):
         self.set_cache({'total_results': 4, 'all_results': self.CACHED_VARIANTS})
         variants, total = query_variants(self.results_model, user=self.user)
@@ -624,6 +652,28 @@ def _assert_expected_get_single_variant_call(self, mock_call, variant_id, expect
     def test_query_variants(self, mock_call):
         super().test_query_variants(mock_call)
 
+    def _test_exclude_previous_search(self, mock_get_variants, *args, **kwargs):
+        super()._test_exclude_previous_search(
+            mock_get_variants, *args, **kwargs,
+            exclude_key_pairs={}, exclude_keys={'MITO': [1000, 2000]}, num_searches=2,
+        )
+        call_args, call_kwargs, _ = self._get_expected_search_call(*args, inheritance_mode='de_novo', sort=None, num_results=100)
+        mock_get_variants.assert_has_calls([mock.call(*call_args, **call_kwargs)])
+
+        # Test when previous results are cached
+        mock_get_variants.reset_mock()
+        self.mock_redis.get.side_effect = [
+            None,
+            json.dumps({'all_results': [VARIANT1, [VARIANT1, VARIANT2], [VARIANT1, SV_VARIANT1], VARIANT2, [VARIANT4, VARIANT3], SV_VARIANT1]}),
+        ]
+        self.mock_redis.keys.side_effect = [[], ['search_results__abc1234__gnomad']]
+        query_variants(self.results_model, user=self.user)
+        super()._test_exclude_previous_search(
+            mock_get_variants, *args, **kwargs, num_searches=1,
+            exclude_key_pairs={'SNV_INDEL': [[1, 2], [3, 4]], 'SNV_INDEL,SV_WGS': [[1, 12]]},
+            exclude_keys={'SNV_INDEL': [1, 2], 'SV_WGS': [12]},
+        )
+
     def test_cached_query_variants(self):
         Project.objects.filter(id=1).update(genome_version='38')
         super().test_cached_query_variants()

seqr/utils/search/utils.py

@@ -7,7 +7,7 @@
 from clickhouse_search.search import get_clickhouse_variants, format_clickhouse_results, \
     get_clickhouse_cache_results, clickhouse_variant_lookup, get_clickhouse_variant_by_id
 from reference_data.models import GENOME_VERSION_GRCh38, GENOME_VERSION_GRCh37
-from seqr.models import Sample, Individual, Project
+from seqr.models import Sample, Individual, Project, VariantSearchResults
 from seqr.utils.logging_utils import SeqrLogger
 from seqr.utils.redis_utils import safe_redis_get_json, safe_redis_get_wildcard_json, safe_redis_set_json
 from seqr.utils.search.constants import XPOS_SORT_KEY, PRIORITIZED_GENE_SORT, RECESSIVE, COMPOUND_HET, \
@@ -293,6 +293,13 @@ def _query_variants(search_model, user, previous_search_results, genome_version,
         if duplicates:
             raise InvalidSearchException(f'ClinVar pathogenicity {", ".join(sorted(duplicates))} is both included and excluded')
 
+    exclude.pop('previousSearch', None)
+    exclude_previous_hash = exclude.pop('previousSearchHash', None)
+    if exclude_previous_hash:
+        parsed_search.update(backend_specific_call(
+            lambda *args: {}, _get_clickhouse_exclude_keys,
+        )(exclude_previous_hash, user, genome_version))
+
     for annotation_key in ['annotations', 'annotations_secondary']:
         if parsed_search.get(annotation_key):
             parsed_search[annotation_key] = {k: v for k, v in parsed_search[annotation_key].items() if v}
@@ -323,6 +330,34 @@ def _query_variants(search_model, user, previous_search_results, genome_version,
     return variant_results, previous_search_results.get('total_results')
 
 
+def _get_clickhouse_exclude_keys(search_hash, user, genome_version):
+    previous_search_model = VariantSearchResults.objects.get(search_hash=search_hash)
+    cached_results = _get_any_sort_cached_results(previous_search_model)
+    if not cached_results:
+        cached_results = {}
+        _query_variants(previous_search_model, user, cached_results, genome_version)
+    results = cached_results['all_results']
+    exclude_keys = defaultdict(list)
+    exclude_key_pairs = defaultdict(list)
+    for variant in results:
+        if isinstance(variant, list):
+            dt1= variant_dataset_type(variant[0])
+            dt2 = variant_dataset_type(variant[1])
+            dataset_type = dt1 if dt1 == dt2 else ','.join(sorted([dt1, dt2]))
+            exclude_key_pairs[dataset_type].append(sorted([variant[0]['key'], variant[1]['key']]))
+        else:
+            dataset_type = variant_dataset_type(variant)
+            exclude_keys[dataset_type].append(variant['key'])
+    return {'exclude_keys': dict(exclude_keys), 'exclude_key_pairs': dict(exclude_key_pairs)}
+
+
+def variant_dataset_type(variant):
+    if not parse_variant_id(variant['variantId']):
+        sample_type = Sample.SAMPLE_TYPE_WGS if 'endChrom' in variant else Sample.SAMPLE_TYPE_WES
+        return f'{Sample.DATASET_TYPE_SV_CALLS}_{sample_type}'
+    return Sample.DATASET_TYPE_MITO_CALLS if 'mitomapPathogenic' in variant else Sample.DATASET_TYPE_VARIANT_CALLS
+
+
 def get_variant_query_gene_counts(search_model, user):
     return backend_specific_call(
         _get_es_variant_query_gene_counts,

seqr/views/apis/variant_search_api.py

@@ -113,6 +113,8 @@ def _get_or_create_results_model(search_hash, search_context, user):
             ).distinct()
 
         search_dict = search_context.get('search', {})
+        if search_context.get('previousSearchHash') and (search_dict.get('exclude') or {}).get('previousSearch'):
+            search_dict['exclude']['previousSearchHash'] = search_context['previousSearchHash']
         search_model = VariantSearch.objects.filter(search=search_dict).filter(
             Q(created_by=user) | Q(name__isnull=False)).first()
         if not search_model: