Geosearch wrapper (#1073)

+ Add wrapper class for SimpleGeoSearcher
castorini · Mar 8, 2022 · 5773500 · 5773500
1 parent e4c8dbe
commit 5773500
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 0 deletions.
diff --git a/pyserini/search/lucene/__init__.py b/pyserini/search/lucene/__init__.py
@@ -16,6 +16,7 @@
 
 from ._base import JQuery, JDisjunctionMaxQueryGenerator, get_topics, \
     get_topics_with_reader, get_qrels_file, get_qrels
+from ._geo_searcher import LuceneGeoSearcher
 from ._impact_searcher import JImpactSearcherResult, LuceneImpactSearcher
 from ._searcher import JLuceneSearcherResult, LuceneSimilarities, \
     LuceneFusionSearcher, LuceneSearcher
@@ -25,6 +26,7 @@
            'JLuceneSearcherResult',
            'JQuery',
            'LuceneFusionSearcher',
+           'LuceneGeoSearcher',
            'LuceneImpactSearcher',
            'LuceneSearcher',
            'LuceneSimilarities',

diff --git a/pyserini/search/lucene/_geo_searcher.py b/pyserini/search/lucene/_geo_searcher.py
@@ -0,0 +1,81 @@
+#
+# Pyserini: Reproducible IR research with sparse and dense representations
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+This module provides Pyserini's Python search interface to Anserini. The main entry point is the ``LuceneGeoSearcher``
+class, which wraps the Java class ``SimpleGeoSearcher`` in Anserini.
+"""
+
+import logging
+from typing import List
+
+from pyserini.pyclass import autoclass
+from pyserini.search.lucene._base import JQuery
+
+
+logger = logging.getLogger(__name__)
+
+
+# Wrappers around Lucene classes
+JSort = autoclass('org.apache.lucene.search.Sort')
+JLatLonDocValuesField = autoclass('org.apache.lucene.document.LatLonDocValuesField')
+JLatLonShape = autoclass('org.apache.lucene.document.LatLonShape')
+JQueryRelation = autoclass('org.apache.lucene.document.ShapeField$QueryRelation')
+
+# Wrappers around Anserini classes
+JGeoSearcher = autoclass('io.anserini.search.SimpleGeoSearcher')
+JGeoSearcherResult = autoclass('io.anserini.search.SimpleSearcher$Result')
+
+
+class LuceneGeoSearcher:
+    """Wrapper class for ``SimpleGeoSearcher`` in Anserini.
+
+    Parameters
+    ----------
+    index_dir : str
+        Path to Lucene index directory.
+    """
+
+    def __init__(self, index_dir: str):
+        self.index_dir = index_dir
+        self.object = JGeoSearcher(index_dir)
+
+    def search(self, q: JQuery, k: int = 10, sort: JSort = None) -> List[JGeoSearcherResult]:
+        """Search the collection.
+
+        Parameters
+        ----------
+        q : JQuery
+            Lucene query.
+        k : int
+            Number of hits to return.
+        sort : JSort
+            Optional distance sort that allows searcher to return results based on distance to a point.
+
+        Returns
+        -------
+        List[JGeoSearcherResult]
+            List of search results.
+        """
+        if sort:
+            hits = self.object.searchGeo(q, k, sort)
+        else:
+            hits = self.object.searchGeo(q, k)
+        return hits
+
+    def close(self):
+        """Close the searcher."""
+        self.object.close()