Vector similarity search support (#1986)

Author: chayim

redis/commands/search/commands.py

@@ -1,5 +1,6 @@
 import itertools
 import time
+from typing import Dict, Union
 
 from ..helpers import parse_to_dict
 from ._util import to_string
@@ -377,7 +378,17 @@ def info(self):
         it = map(to_string, res)
         return dict(zip(it, it))
 
-    def _mk_query_args(self, query):
+    def get_params_args(self, query_params: Dict[str, Union[str, int, float]]):
+        args = []
+        if len(query_params) > 0:
+            args.append("params")
+            args.append(len(query_params) * 2)
+            for key, value in query_params.items():
+                args.append(key)
+                args.append(value)
+        return args
+
+    def _mk_query_args(self, query, query_params: Dict[str, Union[str, int, float]]):
         args = [self.index_name]
 
         if isinstance(query, str):
@@ -387,9 +398,16 @@ def _mk_query_args(self, query):
             raise ValueError(f"Bad query type {type(query)}")
 
         args += query.get_args()
+        if query_params is not None:
+            args += self.get_params_args(query_params)
+
         return args, query
 
-    def search(self, query):
+    def search(
+        self,
+        query: Union[str, Query],
+        query_params: Dict[str, Union[str, int, float]] = None,
+    ):
         """
         Search the index for a given query, and return a result of documents
 
@@ -401,7 +419,7 @@ def search(self, query):
 
         For more information: https://oss.redis.com/redisearch/Commands/#ftsearch
         """  # noqa
-        args, query = self._mk_query_args(query)
+        args, query = self._mk_query_args(query, query_params=query_params)
         st = time.time()
         res = self.execute_command(SEARCH_CMD, *args)
 
@@ -413,18 +431,26 @@ def search(self, query):
             with_scores=query._with_scores,
         )
 
-    def explain(self, query):
+    def explain(
+        self,
+        query=Union[str, Query],
+        query_params: Dict[str, Union[str, int, float]] = None,
+    ):
         """Returns the execution plan for a complex query.
 
         For more information: https://oss.redis.com/redisearch/Commands/#ftexplain
         """  # noqa
-        args, query_text = self._mk_query_args(query)
+        args, query_text = self._mk_query_args(query, query_params=query_params)
         return self.execute_command(EXPLAIN_CMD, *args)
 
-    def explain_cli(self, query):  # noqa
+    def explain_cli(self, query: Union[str, Query]):  # noqa
         raise NotImplementedError("EXPLAINCLI will not be implemented.")
 
-    def aggregate(self, query):
+    def aggregate(
+        self,
+        query: Union[str, Query],
+        query_params: Dict[str, Union[str, int, float]] = None,
+    ):
         """
         Issue an aggregation query.
 
@@ -445,6 +471,8 @@ def aggregate(self, query):
             cmd = [CURSOR_CMD, "READ", self.index_name] + query.build_args()
         else:
             raise ValueError("Bad query", query)
+        if query_params is not None:
+            cmd += self.get_params_args(query_params)
 
         raw = self.execute_command(*cmd)
         return self._get_aggregate_result(raw, query, has_cursor)

tests/test_search.py

@@ -1521,3 +1521,56 @@ def test_profile_limited(client):
     )
     assert det["Iterators profile"]["Type"] == "INTERSECT"
     assert len(res.docs) == 3  # check also the search result
+
+
+@pytest.mark.redismod
+def test_text_params(modclient):
+    modclient.flushdb()
+    modclient.ft().create_index((TextField("name"),))
+
+    modclient.ft().add_document("doc1", name="Alice")
+    modclient.ft().add_document("doc2", name="Bob")
+    modclient.ft().add_document("doc3", name="Carol")
+
+    params_dict = {"name1": "Alice", "name2": "Bob"}
+    q = Query("@name:($name1 | $name2 )")
+    res = modclient.ft().search(q, query_params=params_dict)
+    assert 2 == res.total
+    assert "doc1" == res.docs[0].id
+    assert "doc2" == res.docs[1].id
+
+
+@pytest.mark.redismod
+def test_numeric_params(modclient):
+    modclient.flushdb()
+    modclient.ft().create_index((NumericField("numval"),))
+
+    modclient.ft().add_document("doc1", numval=101)
+    modclient.ft().add_document("doc2", numval=102)
+    modclient.ft().add_document("doc3", numval=103)
+
+    params_dict = {"min": 101, "max": 102}
+    q = Query("@numval:[$min $max]")
+    res = modclient.ft().search(q, query_params=params_dict)
+
+    assert 2 == res.total
+    assert "doc1" == res.docs[0].id
+    assert "doc2" == res.docs[1].id
+
+
+@pytest.mark.redismod
+def test_geo_params(modclient):
+
+    modclient.flushdb()
+    modclient.ft().create_index((GeoField("g")))
+    modclient.ft().add_document("doc1", g="29.69465, 34.95126")
+    modclient.ft().add_document("doc2", g="29.69350, 34.94737")
+    modclient.ft().add_document("doc3", g="29.68746, 34.94882")
+
+    params_dict = {"lat": "34.95126", "lon": "29.69465", "radius": 1000, "units": "km"}
+    q = Query("@g:[$lon $lat $radius $units]")
+    res = modclient.ft().search(q, query_params=params_dict)
+    assert 3 == res.total
+    assert "doc1" == res.docs[0].id
+    assert "doc2" == res.docs[1].id
+    assert "doc3" == res.docs[2].id