<commit_before>from collections import defaultdict
import logging

from seqr.utils.elasticsearch.constants import MAX_COMPOUND_HET_GENES, HAS_ALT_FIELD_KEYS
from seqr.utils.elasticsearch.es_search import EsSearch

logger = logging.getLogger(__name__)


class EsGeneAggSearch(EsSearch):
    AGGREGATION_NAME = 'gene aggregation'
    CACHED_COUNTS_KEY = None

    def aggregate_by_gene(self):
        searches = [self._search]
        for index_searches in self._index_searches.values():
            searches += [index_search for index_search in index_searches if index_search not in searches]

        for search in searches:
            agg = search.aggs.bucket(
                'genes', 'terms', field='mainTranscript_gene_id', size=MAX_COMPOUND_HET_GENES+1
            )
            if self._no_sample_filters:
                for key in HAS_ALT_FIELD_KEYS:
                    agg.bucket(key, 'terms', field=key, size=10000)
            else:
                agg.metric(
                    'vars_by_gene', 'top_hits', size=100, _source='none'
                )

    def _should_execute_single_search(self, page=1, num_results=100):
        return self._is_single_search(), {}

    def _process_single_search_response(self, parsed_response, page=1, num_results=100, deduplicate=False, **kwargs):
        gene_aggs = {gene_id: {k: counts[k] for k in ['total', 'families']} for gene_id, counts in parsed_response.items()}
        self._add_compound_hets(gene_aggs)

        self.previous_search_results['gene_aggs'] = gene_aggs

        return gene_aggs
<commit_msg>fix gene agg search bug<commit_after>from collections import defaultdict
import logging

from seqr.utils.elasticsearch.constants import MAX_COMPOUND_HET_GENES, HAS_ALT_FIELD_KEYS
from seqr.utils.elasticsearch.es_search import EsSearch

logger = logging.getLogger(__name__)


class EsGeneAggSearch(EsSearch):
    AGGREGATION_NAME = 'gene aggregation'
    CACHED_COUNTS_KEY = None

    def aggregate_by_gene(self):
        searches = [self._search]
        for index_searches in self._index_searches.values():
            searches += [index_search for index_search in index_searches]

        for search in searches:
            agg = search.aggs.bucket(
                'genes', 'terms', field='mainTranscript_gene_id', size=MAX_COMPOUND_HET_GENES+1
            )
            if self._no_sample_filters:
                for key in HAS_ALT_FIELD_KEYS:
                    agg.bucket(key, 'terms', field=key, size=10000)
            else:
                agg.metric(
                    'vars_by_gene', 'top_hits', size=100, _source='none'
                )

    def _should_execute_single_search(self, page=1, num_results=100):
        return self._is_single_search(), {}

    def _process_single_search_response(self, parsed_response, page=1, num_results=100, deduplicate=False, **kwargs):
        gene_aggs = {gene_id: {k: counts[k] for k in ['total', 'families']} for gene_id, counts in parsed_response.items()}
        self._add_compound_hets(gene_aggs)

        self.previous_search_results['gene_aggs'] = gene_aggs

        return gene_aggs

{'fraction_non_alphanumeric': 0.08868894601542417, 'fraction_numerical': 0.010925449871465296, 'mean_word_length': 3.001285347043702, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2758734', 'n_tokens_mistral': 1038, 'n_tokens_neox': 947, 'n_words': 201}