debeir.engines.elasticsearch.executor

  1from typing import Dict, List, Optional, Union
  2
  3import tqdm.asyncio
  4from debeir.core.config import apply_config
  5from debeir.core.document import document_factory
  6from debeir.core.query import GenericElasticsearchQuery
  7from debeir.rankers.transformer_sent_encoder import Encoder
  8from debeir.utils.utils import unpack_coroutine
  9from elasticsearch import AsyncElasticsearch as Elasticsearch
 10
 11
 12class ElasticsearchExecutor:
 13    """
 14    Executes an elasticsearch query given the query generated from the config, topics and query class object.
 15
 16    Computes regular patterns of queries expected from general IR topics and indexes.
 17    Includes:
 18        1. Reranking
 19        2. End-to-End Neural IR
 20        3. Statistical keyword matching
 21    """
 22
 23    def __init__(
 24            self,
 25            topics: Dict[Union[str, int], Dict[str, str]],
 26            client: Elasticsearch,
 27            index_name: str,
 28            output_file: str,
 29            query: GenericElasticsearchQuery,
 30            encoder: Optional[Encoder],
 31            return_size: int = 1000,
 32            test=False,
 33            return_id_only=True,
 34            config=None,
 35    ):
 36        self.topics = {"1": topics["1"]} if test else topics
 37        self.client = client
 38        self.index_name = index_name
 39        self.output_file = output_file
 40        self.return_size = return_size
 41        self.query = query
 42        self.encoder = encoder
 43        self.return_id_only = return_id_only
 44        self.config = config
 45        self.document_cls = document_factory['elasticsearch']
 46
 47    def generate_query(self, topic_num):
 48        """
 49        Generates a query given a topic number from the list of topics
 50
 51        :param topic_num:
 52        """
 53        raise NotImplementedError
 54
 55    def execute_query(self, *args, **kwargs):
 56        """
 57        Execute a query given parameters
 58
 59        :param args:
 60        :param kwargs:
 61        """
 62        raise NotImplementedError
 63
 64    @apply_config
 65    def _update_kwargs(self, **kwargs):
 66        return kwargs
 67
 68    async def run_all_queries(
 69            self, query_type=None, return_results=False,
 70            return_size: int = None, return_id_only: bool = False, **kwargs
 71    ) -> List:
 72        """
 73        A generic function that will asynchronously run all topics using the execute_query() method
 74
 75        :param query_type: Which query to execute. Query_type determines which method is used to generate the queries
 76               from self.query.query_funcs: Dict[str, func]
 77        :param return_results: Whether to return raw results from the client. Useful for analysing results directly or
 78               for computing the BM25 scores for log normalization in NIR-style scoring
 79        :param return_size: Number of documents to return. Overrides the config value if exists.
 80        :param return_id_only: Return the ID of the document only, rather than the full source document.
 81        :param args: Arguments to pass to the execute_query method
 82        :param kwargs: Keyword arguments to pass to the execute_query method
 83        :return:
 84            A list of results if return_results = True else an empty list is returned.
 85        """
 86        if not await self.client.ping():
 87            await self.client.close()
 88            raise RuntimeError(
 89                f"Elasticsearch instance cannot be reached at {self.client}"
 90            )
 91
 92        kwargs = self._update_kwargs(**kwargs)
 93
 94        if return_size is None:
 95            return_size = self.return_size
 96
 97        if return_id_only is None:
 98            return_id_only = self.return_id_only
 99
100        if query_type is None:
101            query_type = self.config.query_type
102
103        kwargs.pop('return_size', None)
104        kwargs.pop('return_id_only', None)
105        kwargs.pop('query_type', None)
106
107        tasks = [
108            self.execute_query(
109                topic_num=topic_num,
110                query_type=query_type,
111                return_size=return_size,
112                return_id_only=return_id_only,
113                **kwargs
114            )
115            for topic_num in self.topics
116        ]
117
118        results = []
119
120        for f in tqdm.asyncio.tqdm.as_completed(tasks, desc="Running Queries"):
121            res = await unpack_coroutine(f)
122
123            if return_results:
124                results.append(res)
125
126        return results