Repository: redis/redis-py
Pull request #2041 (2022-03-09 15:50:15)

Support for Vector Fields for Vector Similarity Search
### Pull Request check-list

_Please make sure to review and check all of these items:_

- [x] Does `$ tox` pass with this change (including linting)?
- [x] Do the CI tests pass with this change (enable it first in your forked repo and wait for the github action build to finish)?
- [x] Is the new or changed code fully tested?
- [x] Is a documentation update included (if this change modifies existing APIs, or introduces new ones)?
- [ ] Is there an example added to the examples folder (if applicable)?
- [ ] Was the change added to CHANGES file?

_NOTE: these things are not required to open a PR and can be done
afterwards / while the PR is open._

### Description of change

_Please provide a description of the change here._


codecov-commenter: # [Codecov](https://codecov.io/gh/redis/redis-py/pull/2041?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) Report
> Merging [#2041](https://codecov.io/gh/redis/redis-py/pull/2041?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) (e2bea34) into [master](https://codecov.io/gh/redis/redis-py/commit/6df00190f386f50c5ac5c0b34f30c547dfd40b5c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) (6df0019) will **decrease** coverage by `10.24%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/redis/redis-py/pull/2041/graphs/tree.svg?width=650&height=150&src=pr&token=yenl5fzxxr&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis)](https://codecov.io/gh/redis/redis-py/pull/2041?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis)

```diff
@@             Coverage Diff             @@
##           master    #2041       +/-   ##
===========================================
- Coverage   92.17%   81.93%   -10.25%     
===========================================
  Files         100      100               
  Lines       21024    20239      -785     
===========================================
- Hits        19379    16582     -2797     
- Misses       1645     3657     +2012     
```


| [Impacted Files](https://codecov.io/gh/redis/redis-py/pull/2041?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) | Coverage Δ | |
|---|---|---|
| [redis/commands/search/field.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvY29tbWFuZHMvc2VhcmNoL2ZpZWxkLnB5) | `100.00% <100.00%> (ø)` | |
| [tests/test\_search.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9zZWFyY2gucHk=) | `98.85% <100.00%> (-0.28%)` | :arrow_down: |
| [tests/test\_cluster.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9jbHVzdGVyLnB5) | `14.75% <0.00%> (-83.67%)` | :arrow_down: |
| [redis/cluster.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvY2x1c3Rlci5weQ==) | `22.03% <0.00%> (-69.90%)` | :arrow_down: |
| [redis/crc.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvY3JjLnB5) | `40.00% <0.00%> (-60.00%)` | :arrow_down: |
| [redis/commands/cluster.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvY29tbWFuZHMvY2x1c3Rlci5weQ==) | `35.04% <0.00%> (-59.83%)` | :arrow_down: |
| [redis/utils.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvdXRpbHMucHk=) | `74.28% <0.00%> (-14.61%)` | :arrow_down: |
| [redis/exceptions.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-cmVkaXMvZXhjZXB0aW9ucy5weQ==) | `86.56% <0.00%> (-11.95%)` | :arrow_down: |
| [tests/conftest.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvY29uZnRlc3QucHk=) | `77.31% <0.00%> (-10.29%)` | :arrow_down: |
| [tests/test\_ssl.py](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis#diff-dGVzdHMvdGVzdF9zc2wucHk=) | `40.80% <0.00%> (-6.22%)` | :arrow_down: |
| ... and [54 more](https://codecov.io/gh/redis/redis-py/pull/2041/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/redis/redis-py/pull/2041?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/redis/redis-py/pull/2041?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis). Last update [6df0019...e2bea34](https://codecov.io/gh/redis/redis-py/pull/2041?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=redis).

dvora-h: Can you add your example to examples.rst? (and test it locally)

Patch:
diff --git a/CHANGES b/CHANGES
index ca2d2e81..95843ec7 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,8 +1,11 @@
 
     * Add `items` parameter to `hset` signature
-  	* Create codeql-analysis.yml (#1988). Thanks @chayim
+    * Create codeql-analysis.yml (#1988). Thanks @chayim
     * Add limited support for Lua scripting with RedisCluster
     * Implement `.lock()` method on RedisCluster
+    * Fix cursor returned by SCAN for RedisCluster & change default target to PRIMARIES
+    * Fix scan_iter for RedisCluster
+    * Remove verbose logging when initializing ClusterPubSub, ClusterPipeline or RedisCluster
 
 * 4.1.3 (Feb 8, 2022)
     * Fix flushdb and flushall (#1926)
diff --git a/redis/cluster.py b/redis/cluster.py
index 09c9ab7d..221df856 100644
--- a/redis/cluster.py
+++ b/redis/cluster.py
@@ -7,7 +7,7 @@ import threading
 import time
 from collections import OrderedDict
 
-from redis.client import CaseInsensitiveDict, PubSub, Redis
+from redis.client import CaseInsensitiveDict, PubSub, Redis, parse_scan
 from redis.commands import CommandsParser, RedisClusterCommands
 from redis.connection import ConnectionPool, DefaultParser, Encoder, parse_url
 from redis.crc import REDIS_CLUSTER_HASH_SLOTS, key_slot
@@ -51,10 +51,14 @@ def get_connection(redis_node, *args, **options):
 
 
 def parse_scan_result(command, res, **options):
-    keys_list = []
-    for primary_res in res.values():
-        keys_list += primary_res[1]
-    return 0, keys_list
+    cursors = {}
+    ret = []
+    for node_name, response in res.items():
+        cursor, r = parse_scan(response, **options)
+        cursors[node_name] = cursor
+        ret += r
+
+    return cursors, ret
 
 
 def parse_pubsub_numsub(command, res, **options):
@@ -244,7 +248,6 @@ class RedisCluster(RedisClusterCommands):
                 "INFO",
                 "SHUTDOWN",
                 "KEYS",
-                "SCAN",
                 "DBSIZE",
                 "BGSAVE",
                 "SLOWLOG GET",
@@ -298,6 +301,7 @@ class RedisCluster(RedisClusterCommands):
                 "FUNCTION LIST",
                 "FUNCTION LOAD",
                 "FUNCTION RESTORE",
+                "SCAN",
                 "SCRIPT EXISTS",
                 "SCRIPT FLUSH",
                 "SCRIPT LOAD",
@@ -518,8 +522,6 @@ class RedisCluster(RedisClusterCommands):
              RedisClusterException:
                  - db (Redis do not support database SELECT in cluster mode)
         """
-        log.info("Creating a new instance of RedisCluster client")
-
         if startup_nodes is None:
             startup_nodes = []
 
@@ -1670,7 +1672,6 @@ class ClusterPubSub(PubSub):
         :type host: str
         :type port: int
         """
-        log.info("Creating new instance of ClusterPubSub")
         self.node = None
         self.set_pubsub_node(redis_cluster, node, host, port)
         connection_pool = (
@@ -1802,7 +1803,6 @@ class ClusterPipeline(RedisCluster):
         **kwargs,
     ):
         """ """
-        log.info("Creating new instance of ClusterPipeline")
         self.command_stack = []
         self.nodes_manager = nodes_manager
         self.commands_parser = commands_parser
diff --git a/redis/commands/cluster.py b/redis/commands/cluster.py
index d68f93bc..bcbfbd59 100644
--- a/redis/commands/cluster.py
+++ b/redis/commands/cluster.py
@@ -1,5 +1,8 @@
+from typing import Iterator, Union
+
 from redis.crc import key_slot
 from redis.exceptions import RedisClusterException, RedisError
+from redis.typing import PatternT
 
 from .core import (
     ACLCommands,
@@ -206,6 +209,41 @@ class ClusterDataAccessCommands(DataAccessCommands):
             **kwargs,
         )
 
+    def scan_iter(
+        self,
+        match: Union[PatternT, None] = None,
+        count: Union[int, None] = None,
+        _type: Union[str, None] = None,
+        **kwargs,
+    ) -> Iterator:
+        # Do the first query with cursor=0 for all nodes
+        cursors, data = self.scan(match=match, count=count, _type=_type, **kwargs)
+        yield from data
+
+        cursors = {name: cursor for name, cursor in cursors.items() if cursor != 0}
+        if cursors:
+            # Get nodes by name
+            nodes = {name: self.get_node(node_name=name) for name in cursors.keys()}
+
+            # Iterate over each node till its cursor is 0
+            kwargs.pop("target_nodes", None)
+            while cursors:
+                for name, cursor in cursors.items():
+                    cur, data = self.scan(
+                        cursor=cursor,
+                        match=match,
+                        count=count,
+                        _type=_type,
+                        target_nodes=nodes[name],
+                        **kwargs,
+                    )
+                    yield from data
+                    cursors[name] = cur[name]
+
+                cursors = {
+                    name: cursor for name, cursor in cursors.items() if cursor != 0
+                }
+
 
 class RedisClusterCommands(
     ClusterMultiKeyCommands,
diff --git a/redis/commands/search/field.py b/redis/commands/search/field.py
index 69e39083..da6667df 100644
--- a/redis/commands/search/field.py
+++ b/redis/commands/search/field.py
@@ -1,3 +1,8 @@
+from typing import List
+
+from redis import DataError
+
+
 class Field:
 
     NUMERIC = "NUMERIC"
@@ -5,11 +10,21 @@ class Field:
     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,
+        )


Test patch:
diff --git a/tests/test_cluster.py b/tests/test_cluster.py
index e9242e00..70cbe4f9 100644
--- a/tests/test_cluster.py
+++ b/tests/test_cluster.py
@@ -1773,29 +1773,60 @@ class TestClusterRedisCommands:
         r.set("a", 1)
         r.set("b", 2)
         r.set("c", 3)
-        cursor, keys = r.scan(target_nodes="primaries")
-        assert cursor == 0
-        assert set(keys) == {b"a", b"b", b"c"}
-        _, keys = r.scan(match="a", target_nodes="primaries")
-        assert set(keys) == {b"a"}
+
+        for target_nodes, nodes in zip(
+            ["primaries", "replicas"], [r.get_primaries(), r.get_replicas()]
+        ):
+            cursors, keys = r.scan(target_nodes=target_nodes)
+            assert sorted(keys) == [b"a", b"b", b"c"]
+            assert sorted(cursors.keys()) == sorted(node.name for node in nodes)
+            assert all(cursor == 0 for cursor in cursors.values())
+
+            cursors, keys = r.scan(match="a*", target_nodes=target_nodes)
+            assert sorted(keys) == [b"a"]
+            assert sorted(cursors.keys()) == sorted(node.name for node in nodes)
+            assert all(cursor == 0 for cursor in cursors.values())
 
     @skip_if_server_version_lt("6.0.0")
     def test_cluster_scan_type(self, r):
         r.sadd("a-set", 1)
+        r.sadd("b-set", 1)
+        r.sadd("c-set", 1)
         r.hset("a-hash", "foo", 2)
         r.lpush("a-list", "aux", 3)
-        _, keys = r.scan(match="a*", _type="SET", target_nodes="primaries")
-        assert set(keys) == {b"a-set"}
+
+        for target_nodes, nodes in zip(
+            ["primaries", "replicas"], [r.get_primaries(), r.get_replicas()]
+        ):
+            cursors, keys = r.scan(_type="SET", target_nodes=target_nodes)
+            assert sorted(keys) == [b"a-set", b"b-set", b"c-set"]
+            assert sorted(cursors.keys()) == sorted(node.name for node in nodes)
+            assert all(cursor == 0 for cursor in cursors.values())
+
+            cursors, keys = r.scan(_type="SET", match="a*", target_nodes=target_nodes)
+            assert sorted(keys) == [b"a-set"]
+            assert sorted(cursors.keys()) == sorted(node.name for node in nodes)
+            assert all(cursor == 0 for cursor in cursors.values())
 
     @skip_if_server_version_lt("2.8.0")
     def test_cluster_scan_iter(self, r):
-        r.set("a", 1)
-        r.set("b", 2)
-        r.set("c", 3)
-        keys = list(r.scan_iter(target_nodes="primaries"))
-        assert set(keys) == {b"a", b"b", b"c"}
-        keys = list(r.scan_iter(match="a", target_nodes="primaries"))
-        assert set(keys) == {b"a"}
+        keys_all = []
+        keys_1 = []
+        for i in range(100):
+            s = str(i)
+            r.set(s, 1)
+            keys_all.append(s.encode("utf-8"))
+            if s.startswith("1"):
+                keys_1.append(s.encode("utf-8"))
+        keys_all.sort()
+        keys_1.sort()
+
+        for target_nodes in ["primaries", "replicas"]:
+            keys = r.scan_iter(target_nodes=target_nodes)
+            assert sorted(keys) == keys_all
+
+            keys = r.scan_iter(match="1*", target_nodes=target_nodes)
+            assert sorted(keys) == keys_1
 
     def test_cluster_randomkey(self, r):
         node = r.get_node_from_key("{foo}")
diff --git a/tests/test_search.py b/tests/test_search.py
index 5dea7398..b94ae058 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -12,7 +12,13 @@ import redis.commands.search.aggregation as aggregations
 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()