Merge branch 'master' into vecsim-example

Author: dvora-h

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)

redis/cluster.py

@@ -7,7 +7,7 @@
 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 @@ def __init__(
              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 @@ def __init__(self, redis_cluster, node=None, host=None, port=None, **kwargs):
         :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 @@ def __init__(
         **kwargs,
     ):
         """ """
-        log.info("Creating new instance of ClusterPipeline")
         self.command_stack = []
         self.nodes_manager = nodes_manager
         self.commands_parser = commands_parser

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 @@ def stralgo(
             **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,

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,
+        )

setup.py

@@ -8,7 +8,7 @@
     long_description_content_type="text/markdown",
     keywords=["Redis", "key-value store", "database"],
     license="MIT",
-    version="4.2.0rc3",
+    version="4.2.0",
     packages=find_packages(
         include=[
             "redis",

tests/test_bloom.py

@@ -34,6 +34,11 @@ def test_create(client):
     assert client.cms().initbydim("cmsDim", 100, 5)
     assert client.cms().initbyprob("cmsProb", 0.01, 0.01)
     assert client.topk().reserve("topk", 5, 100, 5, 0.9)
+
+
+@pytest.mark.redismod
+@pytest.mark.experimental
+def test_tdigest_create(client):
     assert client.tdigest().create("tDigest", 100)
 
 
@@ -306,6 +311,7 @@ def test_topk_incrby(client):
 
 # region Test T-Digest
 @pytest.mark.redismod
+@pytest.mark.experimental
 def test_tdigest_reset(client):
     assert client.tdigest().create("tDigest", 10)
     # reset on empty histogram
@@ -319,6 +325,7 @@ def test_tdigest_reset(client):
 
 
 @pytest.mark.redismod
+@pytest.mark.experimental
 def test_tdigest_merge(client):
     assert client.tdigest().create("to-tDigest", 10)
     assert client.tdigest().create("from-tDigest", 10)
@@ -334,6 +341,7 @@ def test_tdigest_merge(client):
 
 
 @pytest.mark.redismod
+@pytest.mark.experimental
 def test_tdigest_min_and_max(client):
     assert client.tdigest().create("tDigest", 100)
     # insert data-points into sketch
@@ -344,6 +352,7 @@ def test_tdigest_min_and_max(client):
 
 
 @pytest.mark.redismod
+@pytest.mark.experimental
 def test_tdigest_quantile(client):
     assert client.tdigest().create("tDigest", 500)
     # insert data-points into sketch
@@ -359,6 +368,7 @@ def test_tdigest_quantile(client):
 
 
 @pytest.mark.redismod
+@pytest.mark.experimental
 def test_tdigest_cdf(client):
     assert client.tdigest().create("tDigest", 100)
     # insert data-points into sketch

tests/test_cluster.py

@@ -1773,29 +1773,60 @@ def test_cluster_scan(self, r):
         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}")