feat: add with_scores support to search command (#59)

Co-authored-by: Guy Korland <[email protected]>

Author: gkorland

README.md

@@ -55,7 +55,7 @@ res = client.search("search engine")
 res = client.search("search engine", snippet_sizes = {'body': 50})
 
 # Searching with complext parameters:
-q = Query("search engine").verbatim().no_content().paging(0,5)
+q = Query("search engine").verbatim().no_content().with_scores().paging(0,5)
 res = client.search(q)
 
 

redisearch/client.py

@@ -361,7 +361,8 @@ def search(self, query):
         return Result(res,
                       not query._no_content,
                       duration=(time.time() - st) * 1000.0,
-                      has_payload=query._with_payloads)
+                      has_payload=query._with_payloads,
+                      with_scores=query._with_scores)
 
     def explain(self, query):
         args, query_text = self._mk_query_args(query)

redisearch/query.py

@@ -23,6 +23,7 @@ def __init__(self, query_string):
         self._fields = None
         self._verbatim = False
         self._with_payloads = False
+        self._with_scores = False
         self._filters = list()
         self._ids = None
         self._slop = -1
@@ -157,6 +158,9 @@ def get_args(self):
 
         if self._with_payloads:
             args.append('WITHPAYLOADS')
+
+        if self._with_scores:
+            args.append('WITHSCORES')
 
         if self._ids:
             args.append('INKEYS')
@@ -225,6 +229,13 @@ def with_payloads(self):
         """
         self._with_payloads = True
         return self
+
+    def with_scores(self):
+        """
+        Ask the engine to return document search scores
+        """
+        self._with_scores = True
+        return self
 
     def limit_fields(self, *fields):
         """

redisearch/result.py

@@ -8,7 +8,7 @@ class Result(object):
     Represents the result of a search query, and has an array of Document objects
     """
 
-    def __init__(self, res, hascontent, duration=0, has_payload = False):
+    def __init__(self, res, hascontent, duration=0, has_payload = False, with_scores = False):
         """
         - **snippets**: An optional dictionary of the form {field: snippet_size} for snippet formatting
         """
@@ -19,15 +19,20 @@ def __init__(self, res, hascontent, duration=0, has_payload = False):
 
         step = 1
         if hascontent:
-            step = 3 if has_payload else 2
-        else:
-            # we can't have nocontent and payloads in the same response
-            has_payload = False
+            step = step + 1
+        if has_payload:
+            step = step + 1
+        if with_scores:
+            step = step + 1
+
+        offset = 2 if with_scores else 1
 
         for i in xrange(1, len(res), step):
             id = to_string(res[i])
-            payload = to_string(res[i+1]) if has_payload else None
-            fields_offset = 2 if has_payload else 1
+            payload = to_string(res[i+offset]) if has_payload else None
+            #fields_offset = 2 if has_payload else 1
+            fields_offset = offset+1 if has_payload else offset
+            score = float(res[i+1]) if with_scores else None
 
             fields = {} 
             if hascontent:
@@ -40,7 +45,7 @@ def __init__(self, res, hascontent, duration=0, has_payload = False):
             except KeyError:
                 pass
 
-            doc = Document(id, payload=payload, **fields)
+            doc = Document(id, score=score, payload=payload, **fields) if with_scores else Document(id, payload=payload, **fields)
             self.docs.append(doc)
 
     def __repr__(self):

test/test.py

@@ -194,6 +194,30 @@ def testPayloads(self):
             self.assertEqual('foo baz', res.docs[1].payload)
             self.assertIsNone(res.docs[0].payload)
 
+    def testScores(self):
+        
+        conn = self.redis()
+
+        with conn as r:
+            # Creating a client with a given index name
+            client = Client('idx', port=conn.port)
+            client.redis.flushdb()
+            client.create_index((TextField('txt'),))
+
+            client.add_document('doc1', txt = 'foo baz')
+            client.add_document('doc2', txt = 'foo bar')
+
+            q = Query("foo ~bar").with_scores()
+            res = client.search(q)
+            print("RES", res)
+            self.assertEqual(2, res.total)
+            
+            self.assertEqual('doc2', res.docs[0].id)
+            self.assertEqual(3.0, res.docs[0].score)
+
+            self.assertEqual('doc1', res.docs[1].id)
+            self.assertEqual(0.2, res.docs[1].score)
+
     def testReplace(self):
 
         conn = self.redis()