Repository: castorini/pyserini

Issue:
searcher can't take custom constructed query with RM3 set
With RM3, we get an NPE when we try to search with a query build using the querybuilder.



Discussion:

Resolving patch:
diff --git a/pyserini/index/pyutils.py b/pyserini/index/pyutils.py
--- a/pyserini/index/pyutils.py
+++ b/pyserini/index/pyutils.py
@@ -258,7 +258,7 @@ def doc_contents(self, docid: str) -> str:
         """
         return self.object.documentContents(self.reader, JString(docid))
 
-    def compute_bm25_term_weight(self, docid: str, term: str, k1=0.9, b=0.4) -> float:
+    def compute_bm25_term_weight(self, docid: str, term: str, analyzer=get_lucene_analyzer(), k1=0.9, b=0.4) -> float:
         """Computes the BM25 weight of an (analyzed) term in a document. Note that this method takes the analyzed
         (i.e., stemmed) form because the most common use case is to take the term from the output of
         :func:`get_document_vector`.
@@ -279,8 +279,20 @@ def compute_bm25_term_weight(self, docid: str, term: str, k1=0.9, b=0.4) -> floa
         float
             The BM25 weight of the term in the document, or 0 if the term does not exist in the document.
         """
-        return self.object.getBM25TermWeightWithParameters(self.reader, JString(docid), JString(term.encode('utf-8')),
-                                                           float(k1), float(b))
+        if analyzer is None:
+            return self.object.getBM25AnalyzedTermWeightWithParameters(self.reader, JString(docid),
+                                                                       JString(term.encode('utf-8')),
+                                                                       float(k1), float(b))
+        else:
+            return self.object.getBM25UnanalyzedTermWeightWithParameters(self.reader, JString(docid),
+                                                                         JString(term.encode('utf-8')), analyzer,
+                                                                         float(k1), float(b))
+
+    def compute_query_document_score(self, docid: str, query: str, similarity=None):
+        if similarity is None:
+            return self.object.computeQueryDocumentScore(self.reader, docid, query)
+        else:
+            return self.object.computeQueryDocumentScoreWithSimilarity(self.reader, docid, query, similarity)
 
     def convert_internal_docid_to_collection_docid(self, docid: int) -> str:
         """Converts Lucene's internal ``docid`` to its external collection ``docid``.
diff --git a/pyserini/pyclass.py b/pyserini/pyclass.py
--- a/pyserini/pyclass.py
+++ b/pyserini/pyclass.py
@@ -121,4 +121,3 @@ class JCollections(Enum):
     TweetCollection = autoclass('io.anserini.collection.TweetCollection')
     WashingtonPostCollection = autoclass('io.anserini.collection.WashingtonPostCollection')
     WikipediaCollection = autoclass('io.anserini.collection.WikipediaCollection')
-
diff --git a/pyserini/search/pysearch.py b/pyserini/search/pysearch.py
--- a/pyserini/search/pysearch.py
+++ b/pyserini/search/pysearch.py
@@ -23,7 +23,7 @@
 from typing import Dict, List, Union
 
 from ..pyclass import JSimpleSearcher, JSimpleSearcherResult, JDocument, JString, JArrayList, JTopics, JTopicReader, \
-    JQueryGenerator, JSimpleNearestNeighborSearcherResult, JSimpleNearestNeighborSearcher, JQuery
+    JQueryGenerator, JSimpleNearestNeighborSearcherResult, JSimpleNearestNeighborSearcher, JQuery, autoclass
 
 logger = logging.getLogger(__name__)
 
@@ -80,7 +80,7 @@ def search(self, q: Union[str, JQuery], k: int = 10,
         Parameters
         ----------
         q : Union[str, JQuery]
-            The query string / The JQuery.
+            The query string or the ``JQuery`` objected.
         k : int
             The number of hits to return.
         query_generator : JQueryGenerator
@@ -94,6 +94,13 @@ def search(self, q: Union[str, JQuery], k: int = 10,
         if query_generator:
             return self.object.search(query_generator, JString(q), k)
         elif isinstance(q, JQuery):
+            # Note that RM3 requires the notion of a query (string) to estimate the appropriate models. If we're just
+            # given a Lucene query, it's unclear what the "query" is for this estimation. One possibility is to extract
+            # all the query terms from the Lucene query, although this might yield unexpected behavior from the user's
+            # perspective. Until we think through what exactly is the "right thing to do", we'll raise an exception
+            # here explicitly.
+            if self.is_using_rm3():
+                raise NotImplementedError('RM3 incompatible with search using a Lucene query.')
             return self.object.search(q, k)
         else:
             return self.object.search(JString(q.encode('utf8')), k)
@@ -133,11 +140,12 @@ def batch_search(self, queries: List[str], qids: List[str], k: int = 10,
         return {r.getKey(): r.getValue() for r in results}
 
     def search_fields(self, q, f, boost, k):
-        """
+        """Searches the collection, scoring a separate field with a boost weight.
+
         Parameters
         ----------
         q : str
-            Query string
+            Query string.
         f : str
             Name of additional field to search over
         boost : float
@@ -161,59 +169,58 @@ def set_analyzer(self, analyzer):
         """
         self.object.setAnalyzer(analyzer)
 
-    def set_search_tweets(self, flag):
-        """
-        Parameters
-        ----------
-        flag : bool
-            True if searching over tweets
-        """
-        self.object.setSearchTweets(flag)
+    def set_rm3(self, fb_terms=10, fb_docs=10, original_query_weight=float(0.5), rm3_output_query=False):
+        """Configures RM3 query expansion.
 
-    def set_rm3_reranker(self, fb_terms=10, fb_docs=10,
-                         original_query_weight=float(0.5),
-                         rm3_output_query=False):
-        """
         Parameters
         ----------
         fb_terms : int
-            RM3 parameter for number of expansion terms
+            RM3 parameter for number of expansion terms.
         fb_docs : int
-            RM3 parameter for number of documents
+            RM3 parameter for number of expansion documents.
         original_query_weight : float
-            RM3 parameter for weight to assign to the original query
+            RM3 parameter for weight to assign to the original query.
         rm3_output_query : bool
-            True if we want to print original and expanded queries for RM3
+            Whether we want to print the original and expanded as debug output.
         """
-        self.object.setRM3Reranker(fb_terms, fb_docs,
-                                   original_query_weight, rm3_output_query)
+        self.object.setRM3(fb_terms, fb_docs, original_query_weight, rm3_output_query)
 
-    def unset_rm3_reranker(self):
-        """
-        Parameters
-        ----------
+    def unset_rm3(self):
+        """Turns off use of RM3 query expansion.
         """
-        self.object.unsetRM3Reranker()
+        self.object.unsetRM3()
 
-    def set_lm_dirichlet_similarity(self, mu):
+    def is_using_rm3(self) -> bool:
+        """Returns whether or not RM3 query expansion is being performed.
         """
+        return self.object.useRM3()
+
+    def set_qld(self, mu=float(1000)):
+        """Configures query likelihood with Dirichlet smoothing as the scoring function.
+
         Parameters
         ----------
         mu : float
-            Dirichlet smoothing parameter
+            Dirichlet smoothing parameter mu.
         """
-        self.object.setLMDirichletSimilarity(float(mu))
+        self.object.setQLD(float(mu))
+
+    def set_bm25(self, k1=float(0.9), b=float(0.4)):
+        """Configures BM25 as the scoring function.
 
-    def set_bm25_similarity(self, k1, b):
-        """
         Parameters
         ----------
         k1 : float
-            BM25 k1 parameter
+            BM25 k1 parameter.
         b : float
-            BM25 b parameter
+            BM25 b parameter.
         """
-        self.object.setBM25Similarity(float(k1), float(b))
+        self.object.setBM25(float(k1), float(b))
+
+    def get_similarity(self):
+        """Returns the Lucene ``Similarity`` used as the scoring function.
+        """
+        return self.object.getSimilarity()
 
     def doc(self, docid: Union[str, int]) -> Document:
         """Returns the :class:`Document` corresponding to ``docid``. The ``docid`` is overloaded: if it is of type
@@ -305,6 +312,16 @@ def get_topics(collection_name):
     return t
 
 
+class LuceneSimilarities:
+    @staticmethod
+    def bm25(k1=0.9, b=0.4):
+        return autoclass('org.apache.lucene.search.similarities.BM25Similarity')(k1, b)
+
+    @staticmethod
+    def qld(mu=1000):
+        return autoclass('org.apache.lucene.search.similarities.LMDirichletSimilarity')(mu)
+
+
 class SimpleNearestNeighborSearcher:
 
     def __init__(self, index_dir: str):


Test patch:
diff --git a/tests/test_indexutils.py b/tests/test_indexutils.py
--- a/tests/test_indexutils.py
+++ b/tests/test_indexutils.py
@@ -24,6 +24,7 @@
 from pyserini.analysis import pyanalysis
 from pyserini.index import pyutils
 from pyserini.pyclass import JString
+from pyserini.search import pysearch
 
 
 class TestIndexUtils(unittest.TestCase):
@@ -40,7 +41,8 @@ def setUp(self):
         tarball.extractall(self.index_dir)
         tarball.close()
 
-        self.index_utils = pyutils.IndexReaderUtils('{}lucene-index.cacm'.format(self.index_dir))
+        self.searcher = pysearch.SimpleSearcher(f'{self.index_dir}lucene-index.cacm')
+        self.index_utils = pyutils.IndexReaderUtils(f'{self.index_dir}lucene-index.cacm')
 
     def test_terms_count(self):
         # We're going to iterate through the index and make sure we have the correct number of terms.
@@ -180,14 +182,34 @@ def test_doc_by_field(self):
                          self.index_utils.doc_by_field('id', 'CACM-3134').docid())
 
     def test_bm25_weight(self):
-        self.assertAlmostEqual(self.index_utils.compute_bm25_term_weight('CACM-3134', 'inform', k1=1.2, b=0.75),
-                               1.925014, places=5)
-        self.assertAlmostEqual(self.index_utils.compute_bm25_term_weight('CACM-3134', 'retriev', k1=1.2, b=0.75),
-                               2.496352, places=5)
-
-        self.assertAlmostEqual(self.index_utils.compute_bm25_term_weight('CACM-3134', 'inform'), 2.06514, places=5)
-        self.assertAlmostEqual(self.index_utils.compute_bm25_term_weight('CACM-3134', 'retriev'), 2.70038, places=5)
-
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'inform', analyzer=None, k1=1.2, b=0.75),
+            1.925014, places=5)
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'information', k1=1.2, b=0.75),
+            1.925014, places=5)
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'retriev', analyzer=None, k1=1.2, b=0.75),
+            2.496352, places=5)
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'retrieval', k1=1.2, b=0.75),
+            2.496352, places=5)
+
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'inform',  analyzer=None),
+            2.06514, places=5)
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'information'),
+            2.06514, places=5)
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'retriev',  analyzer=None),
+            2.70038, places=5)
+        self.assertAlmostEqual(
+            self.index_utils.compute_bm25_term_weight('CACM-3134', 'retrieval'),
+            2.70038, places=5)
+
+        self.assertAlmostEqual(self.index_utils.compute_bm25_term_weight('CACM-3134', 'fox',  analyzer=None),
+                               0., places=5)
         self.assertAlmostEqual(self.index_utils.compute_bm25_term_weight('CACM-3134', 'fox'), 0., places=5)
 
     def test_docid_converstion(self):
@@ -202,6 +224,46 @@ def test_jstring_term(self):
             # Should fail when pyjnius has solved this internally.
             JString('zoölogy')
 
+    def test_query_doc_score_default(self):
+        queries = ['information retrieval', 'databases']
+
+        for query in queries:
+            hits = self.searcher.search(query)
+
+            # We're going to verify that the score of each hit is about the same as the output of
+            # compute_query_document_score
+            for i in range(0, len(hits)):
+                self.assertAlmostEqual(hits[i].score,
+                                       self.index_utils.compute_query_document_score(hits[i].docid, query), places=4)
+
+    def test_query_doc_score_custom_similarity(self):
+        custom_bm25 = pysearch.LuceneSimilarities.bm25(0.8, 0.2)
+        queries = ['information retrieval', 'databases']
+        self.searcher.set_bm25(0.8, 0.2)
+
+        for query in queries:
+            hits = self.searcher.search(query)
+
+            # We're going to verify that the score of each hit is about the same as the output of
+            # compute_query_document_score
+            for i in range(0, len(hits)):
+                self.assertAlmostEqual(hits[i].score,
+                                       self.index_utils.compute_query_document_score(
+                                           hits[i].docid, query, similarity=custom_bm25), places=4)
+
+        custom_qld = pysearch.LuceneSimilarities.qld(500)
+        self.searcher.set_qld(500)
+
+        for query in queries:
+            hits = self.searcher.search(query)
+
+            # We're going to verify that the score of each hit is about the same as the output of
+            # compute_query_document_score
+            for i in range(0, len(hits)):
+                self.assertAlmostEqual(hits[i].score,
+                                       self.index_utils.compute_query_document_score(
+                                           hits[i].docid, query, similarity=custom_qld), places=4)
+
     def tearDown(self):
         os.remove(self.tarball_name)
         shutil.rmtree(self.index_dir)
diff --git a/tests/test_querybuilding.py b/tests/test_querybuilding.py
--- a/tests/test_querybuilding.py
+++ b/tests/test_querybuilding.py
@@ -82,23 +82,35 @@ def testBuildBoostedQuery(self):
             self.assertNotEqual(h1.score, h3.score)
 
     def testTermQuery(self):
-        term_query1 = pyquerybuilder.get_term_query('information')
-        term_query2 = pyquerybuilder.get_term_query('retrieval')
-
         should = pyquerybuilder.JBooleanClauseOccur['should'].value
+        query_builder = pyquerybuilder.get_boolean_query_builder()
+        query_builder.add(pyquerybuilder.get_term_query('information'), should)
+        query_builder.add(pyquerybuilder.get_term_query('retrieval'), should)
 
-        boolean_query1 = pyquerybuilder.get_boolean_query_builder()
-        boolean_query1.add(term_query1, should)
-        boolean_query1.add(term_query2, should)
-
-        bq1 = boolean_query1.build()
-        hits1 = self.searcher.search(bq1)
+        query = query_builder.build()
+        hits1 = self.searcher.search(query)
         hits2 = self.searcher.search('information retrieval')
 
         for h1, h2 in zip(hits1, hits2):
             self.assertEqual(h1.docid, h2.docid)
             self.assertEqual(h1.score, h2.score)
 
+    def testIncompatabilityWithRM3(self):
+        should = pyquerybuilder.JBooleanClauseOccur['should'].value
+        query_builder = pyquerybuilder.get_boolean_query_builder()
+        query_builder.add(pyquerybuilder.get_term_query('information'), should)
+        query_builder.add(pyquerybuilder.get_term_query('retrieval'), should)
+
+        query = query_builder.build()
+        hits = self.searcher.search(query)
+        self.assertEqual(10, len(hits))
+
+        self.searcher.set_rm3()
+        self.assertTrue(self.searcher.is_using_rm3())
+
+        with self.assertRaises(NotImplementedError):
+            self.searcher.search(query)
+
     def testTermQuery2(self):
         term_query1 = pyquerybuilder.get_term_query('inform', analyzer=get_lucene_analyzer(stemming=False))
         term_query2 = pyquerybuilder.get_term_query('retriev', analyzer=get_lucene_analyzer(stemming=False))
diff --git a/tests/test_search.py b/tests/test_search.py
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -22,7 +22,7 @@
 from typing import List, Dict
 from urllib.request import urlretrieve
 
-from pyserini.pyclass import JSimpleSearcherResult
+from pyserini.pyclass import JSimpleSearcherResult, autoclass
 from pyserini.search import pysearch
 
 
@@ -43,6 +43,8 @@ def setUp(self):
         self.searcher = pysearch.SimpleSearcher(f'{self.index_dir}lucene-index.cacm')
 
     def test_basic(self):
+        self.assertTrue(self.searcher.get_similarity().toString().startswith('BM25'))
+
         hits = self.searcher.search('information retrieval')
 
         self.assertEqual(3204, self.searcher.num_docs)
@@ -119,6 +121,82 @@ def test_batch_k(self):
         self.assertTrue(isinstance(results['q2'][0], JSimpleSearcherResult))
         self.assertEqual(len(results['q2']), 100)
 
+    def test_different_similarity(self):
+        # qld, default mu
+        self.searcher.set_qld()
+        self.assertTrue(self.searcher.get_similarity().toString().startswith('LM Dirichlet'))
+
+        hits = self.searcher.search('information retrieval')
+
+        self.assertEqual(hits[0].docid, 'CACM-3134')
+        self.assertAlmostEqual(hits[0].score, 3.68030, places=5)
+        self.assertEqual(hits[9].docid, 'CACM-1927')
+        self.assertAlmostEqual(hits[9].score, 2.53240, places=5)
+
+        # bm25, default parameters
+        self.searcher.set_bm25()
+        self.assertTrue(self.searcher.get_similarity().toString().startswith('BM25'))
+
+        hits = self.searcher.search('information retrieval')
+
+        self.assertEqual(hits[0].docid, 'CACM-3134')
+        self.assertAlmostEqual(hits[0].score, 4.76550, places=5)
+        self.assertEqual(hits[9].docid, 'CACM-2516')
+        self.assertAlmostEqual(hits[9].score, 4.21740, places=5)
+
+        # qld, custom mu
+        self.searcher.set_qld(100)
+        self.assertTrue(self.searcher.get_similarity().toString().startswith('LM Dirichlet'))
+
+        hits = self.searcher.search('information retrieval')
+
+        self.assertEqual(hits[0].docid, 'CACM-3134')
+        self.assertAlmostEqual(hits[0].score, 6.35580, places=5)
+        self.assertEqual(hits[9].docid, 'CACM-2631')
+        self.assertAlmostEqual(hits[9].score, 5.18960, places=5)
+
+        # bm25, custom parameters
+        self.searcher.set_bm25(0.8, 0.3)
+        self.assertTrue(self.searcher.get_similarity().toString().startswith('BM25'))
+
+        hits = self.searcher.search('information retrieval')
+
+        self.assertEqual(hits[0].docid, 'CACM-3134')
+        self.assertAlmostEqual(hits[0].score, 4.86880, places=5)
+        self.assertEqual(hits[9].docid, 'CACM-2516')
+        self.assertAlmostEqual(hits[9].score, 4.33320, places=5)
+
+    def test_rm3(self):
+        self.searcher.set_rm3()
+        self.assertTrue(self.searcher.is_using_rm3())
+
+        hits = self.searcher.search('information retrieval')
+
+        self.assertEqual(hits[0].docid, 'CACM-3134')
+        self.assertAlmostEqual(hits[0].score, 2.18010, places=5)
+        self.assertEqual(hits[9].docid, 'CACM-2516')
+        self.assertAlmostEqual(hits[9].score, 1.70330, places=5)
+
+        self.searcher.unset_rm3()
+        self.assertFalse(self.searcher.is_using_rm3())
+
+        hits = self.searcher.search('information retrieval')
+
+        self.assertEqual(hits[0].docid, 'CACM-3134')
+        self.assertAlmostEqual(hits[0].score, 4.76550, places=5)
+        self.assertEqual(hits[9].docid, 'CACM-2516')
+        self.assertAlmostEqual(hits[9].score, 4.21740, places=5)
+
+        self.searcher.set_rm3(fb_docs=4, fb_terms=6, original_query_weight=0.3)
+        self.assertTrue(self.searcher.is_using_rm3())
+
+        hits = self.searcher.search('information retrieval')
+
+        self.assertEqual(hits[0].docid, 'CACM-3134')
+        self.assertAlmostEqual(hits[0].score, 2.17190, places=5)
+        self.assertEqual(hits[9].docid, 'CACM-1457')
+        self.assertAlmostEqual(hits[9].score, 1.43700, places=5)
+
     def test_doc_int(self):
         # The doc method is overloaded: if input is int, it's assumed to be a Lucene internal docid.
         doc = self.searcher.doc(1)