Support for Vector Fields for Vector Similarity Search (#2041)

* Support Vector field in FT.CREATE command

* linters

* fix data error

* change to dic

* add type hints and docstring to constructor

* test not supported algorithm

* linters

* fix errors

* example

* delete example

Co-authored-by: dvora-h <[email protected]>

redis/commands/search/field.py

@@ -1,15 +1,30 @@
+from typing import List
+
+from redis import DataError
+
+
 class Field:
 
     NUMERIC = "NUMERIC"
     TEXT = "TEXT"
     WEIGHT = "WEIGHT"
     GEO = "GEO"
     TAG = "TAG"
+    VECTOR = "VECTOR"
     SORTABLE = "SORTABLE"
     NOINDEX = "NOINDEX"
     AS = "AS"
 
-    def __init__(self, name, args=[], sortable=False, no_index=False, as_name=None):
+    def __init__(
+        self,
+        name: str,
+        args: List[str] = None,
+        sortable: bool = False,
+        no_index: bool = False,
+        as_name: str = None,
+    ):
+        if args is None:
+            args = []
         self.name = name
         self.args = args
         self.args_suffix = list()
@@ -44,7 +59,12 @@ class TextField(Field):
     PHONETIC = "PHONETIC"
 
     def __init__(
-        self, name, weight=1.0, no_stem=False, phonetic_matcher=None, **kwargs
+        self,
+        name: str,
+        weight: float = 1.0,
+        no_stem: bool = False,
+        phonetic_matcher: str = None,
+        **kwargs,
     ):
         Field.__init__(self, name, args=[Field.TEXT, Field.WEIGHT, weight], **kwargs)
 
@@ -65,7 +85,7 @@ class NumericField(Field):
     NumericField is used to define a numeric field in a schema definition
     """
 
-    def __init__(self, name, **kwargs):
+    def __init__(self, name: str, **kwargs):
         Field.__init__(self, name, args=[Field.NUMERIC], **kwargs)
 
 
@@ -74,7 +94,7 @@ class GeoField(Field):
     GeoField is used to define a geo-indexing field in a schema definition
     """
 
-    def __init__(self, name, **kwargs):
+    def __init__(self, name: str, **kwargs):
         Field.__init__(self, name, args=[Field.GEO], **kwargs)
 
 
@@ -86,7 +106,52 @@ class TagField(Field):
 
     SEPARATOR = "SEPARATOR"
 
-    def __init__(self, name, separator=",", **kwargs):
+    def __init__(self, name: str, separator: str = ",", **kwargs):
         Field.__init__(
             self, name, args=[Field.TAG, self.SEPARATOR, separator], **kwargs
         )
+
+
+class VectorField(Field):
+    """
+    Allows vector similarity queries against the value in this attribute.
+    See https://oss.redis.com/redisearch/Vectors/#vector_fields.
+    """
+
+    def __init__(self, name: str, algorithm: str, attributes: dict, **kwargs):
+        """
+        Create Vector Field. Notice that Vector cannot have sortable or no_index tag,
+        although it's also a Field.
+
+        ``name`` is the name of the field.
+
+        ``algorithm`` can be "FLAT" or "HNSW".
+
+        ``attributes`` each algorithm can have specific attributes. Some of them
+        are mandatory and some of them are optional. See
+        https://oss.redis.com/redisearch/master/Vectors/#specific_creation_attributes_per_algorithm
+        for more information.
+        """
+        sort = kwargs.get("sortable", False)
+        noindex = kwargs.get("no_index", False)
+
+        if sort or noindex:
+            raise DataError("Cannot set 'sortable' or 'no_index' in Vector fields.")
+
+        if algorithm.upper() not in ["FLAT", "HNSW"]:
+            raise DataError(
+                "Realtime vector indexing supporting 2 Indexing Methods:"
+                "'FLAT' and 'HNSW'."
+            )
+
+        attr_li = []
+
+        for key, value in attributes.items():
+            attr_li.extend([key, value])
+
+        Field.__init__(
+            self,
+            name,
+            args=[Field.VECTOR, algorithm, len(attr_li), *attr_li],
+            **kwargs,
+        )

tests/test_search.py

@@ -12,7 +12,13 @@
 import redis.commands.search.reducers as reducers
 from redis.commands.json.path import Path
 from redis.commands.search import Search
-from redis.commands.search.field import GeoField, NumericField, TagField, TextField
+from redis.commands.search.field import (
+    GeoField,
+    NumericField,
+    TagField,
+    TextField,
+    VectorField,
+)
 from redis.commands.search.indexDefinition import IndexDefinition, IndexType
 from redis.commands.search.query import GeoFilter, NumericFilter, Query
 from redis.commands.search.result import Result
@@ -1522,6 +1528,40 @@ def test_profile_limited(client):
     assert len(res.docs) == 3  # check also the search result
 
 
+@pytest.mark.redismod
+def test_vector_field(modclient):
+    modclient.flushdb()
+    modclient.ft().create_index(
+        (
+            VectorField(
+                "v", "HNSW", {"TYPE": "FLOAT32", "DIM": 2, "DISTANCE_METRIC": "L2"}
+            ),
+        )
+    )
+    modclient.hset("a", "v", "aaaaaaaa")
+    modclient.hset("b", "v", "aaaabaaa")
+    modclient.hset("c", "v", "aaaaabaa")
+
+    q = Query("*=>[KNN 2 @v $vec]").return_field("__v_score").sort_by("__v_score", True)
+    res = modclient.ft().search(q, query_params={"vec": "aaaaaaaa"})
+
+    assert "a" == res.docs[0].id
+    assert "0" == res.docs[0].__getattribute__("__v_score")
+
+
+@pytest.mark.redismod
+def test_vector_field_error(modclient):
+    modclient.flushdb()
+
+    # sortable tag
+    with pytest.raises(Exception):
+        modclient.ft().create_index((VectorField("v", "HNSW", {}, sortable=True),))
+
+    # not supported algorithm
+    with pytest.raises(Exception):
+        modclient.ft().create_index((VectorField("v", "SORT", {}),))
+
+
 @pytest.mark.redismod
 def test_text_params(modclient):
     modclient.flushdb()