Redis Vector Sets support

.gitignore

@@ -8,3 +8,4 @@ results/*
 tools/custom/data.json
 
 *.png
+venv/

engine/clients/client_factory.py

@@ -31,6 +31,12 @@
     WeaviateUploader,
 )
 
+from engine.clients.vectorsets import (
+    RedisVsetConfigurator,
+    RedisVsetSearcher,
+    RedisVsetUploader,
+)
+
 ENGINE_CONFIGURATORS = {
     "qdrant": QdrantConfigurator,
     "weaviate": WeaviateConfigurator,
@@ -39,6 +45,7 @@
     "opensearch": OpenSearchConfigurator,
     "redis": RedisConfigurator,
     "pgvector": PgVectorConfigurator,
+    "vectorsets": RedisVsetConfigurator,
 }
 
 ENGINE_UPLOADERS = {
@@ -49,6 +56,7 @@
     "opensearch": OpenSearchUploader,
     "redis": RedisUploader,
     "pgvector": PgVectorUploader,
+    "vectorsets": RedisVsetUploader,
 }
 
 ENGINE_SEARCHERS = {
@@ -59,6 +67,7 @@
     "opensearch": OpenSearchSearcher,
     "redis": RedisSearcher,
     "pgvector": PgVectorSearcher,
+    "vectorsets": RedisVsetSearcher,
 }
 
 

engine/clients/vectorsets/__init__.py

@@ -0,0 +1,3 @@
+from engine.clients.vectorsets.configure import RedisVsetConfigurator
+from engine.clients.vectorsets.search import RedisVsetSearcher
+from engine.clients.vectorsets.upload import RedisVsetUploader

engine/clients/vectorsets/config.py

@@ -0,0 +1,9 @@
+import os
+
+REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
+REDIS_AUTH = os.getenv("REDIS_AUTH", None)
+REDIS_USER = os.getenv("REDIS_USER", None)
+REDIS_CLUSTER = bool(int(os.getenv("REDIS_CLUSTER", 0)))
+
+# 90 seconds timeout
+REDIS_QUERY_TIMEOUT = int(os.getenv("REDIS_QUERY_TIMEOUT", 90 * 1000))

engine/clients/vectorsets/configure.py

@@ -0,0 +1,44 @@
+import redis
+from redis import Redis, RedisCluster
+
+from benchmark.dataset import Dataset
+from engine.base_client.configure import BaseConfigurator
+from engine.clients.vectorsets.config import (
+    REDIS_AUTH,
+    REDIS_CLUSTER,
+    REDIS_PORT,
+    REDIS_USER,
+)
+
+
+class RedisVsetConfigurator(BaseConfigurator):
+
+    def __init__(self, host, collection_params: dict, connection_params: dict):
+        super().__init__(host, collection_params, connection_params)
+        redis_constructor = RedisCluster if REDIS_CLUSTER else Redis
+        self._is_cluster = True if REDIS_CLUSTER else False
+        self.client = redis_constructor(
+            host=host, port=REDIS_PORT, password=REDIS_AUTH, username=REDIS_USER
+        )
+        self.client.flushall()
+
+    def clean(self):
+        conns = [self.client]
+        if self._is_cluster:
+            conns = [
+                self.client.get_redis_connection(node)
+                for node in self.client.get_primaries()
+            ]
+        for conn in conns:
+            index = conn.ft()
+            try:
+                 conn.flushall()
+            except redis.ResponseError as e:
+                print(e)
+
+    def recreate(self, dataset: Dataset, collection_params):
+        pass
+
+
+if __name__ == "__main__":
+    pass

engine/clients/vectorsets/search.py

@@ -0,0 +1,53 @@
+import random
+from typing import List, Tuple
+
+import numpy as np
+from redis import Redis, RedisCluster
+
+
+from engine.base_client.search import BaseSearcher
+from engine.clients.vectorsets.config import (
+    REDIS_AUTH,
+    REDIS_CLUSTER,
+    REDIS_PORT,
+    REDIS_QUERY_TIMEOUT,
+    REDIS_USER,
+)
+from engine.clients.redis.parser import RedisConditionParser
+
+
+class RedisVsetSearcher(BaseSearcher):
+    search_params = {}
+    client = None
+    parser = RedisConditionParser()
+
+    @classmethod
+    def init_client(cls, host, distance, connection_params: dict, search_params: dict):
+        redis_constructor = RedisCluster if REDIS_CLUSTER else Redis
+        cls.client = redis_constructor(
+            host=host, port=REDIS_PORT, password=REDIS_AUTH, username=REDIS_USER
+        )
+        cls.search_params = search_params
+        cls._is_cluster = True if REDIS_CLUSTER else False
+        # In the case of CLUSTER API enabled we randomly select the starting primary shard
+        # when doing the client initialization to evenly distribute the load among the cluster
+        cls.conns = [cls.client]
+        if cls._is_cluster:
+            cls.conns = [
+                cls.client.get_redis_connection(node)
+                for node in cls.client.get_primaries()
+            ]
+        cls._ft = cls.conns[random.randint(0, len(cls.conns)) - 1].ft()
+
+    @classmethod
+    def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
+        ef = cls.search_params["search_params"]["ef"]
+        response = cls.client.execute_command("VSIM", "idx", "FP32", np.array(vector).astype(np.float32).tobytes(), "WITHSCORES", "COUNT", top, "EF", ef)
+        # decode responses
+        # every even cell is id, every odd is the score
+        # scores needs to be 1 - scores since on vector sets 1 is identical, 0 is opposite vector
+        ids = [int(response[i]) for i in range(0, len(response), 2)]
+        scores = [1 - float(response[i]) for i in range(1, len(response), 2)]
+        # we need to return a list of tuples
+        # where the first element is the id and the second is the score
+        return list(zip(ids, scores))        

engine/clients/vectorsets/upload.py

@@ -0,0 +1,48 @@
+from typing import List, Optional
+
+import numpy as np
+from redis import Redis, RedisCluster
+
+from engine.base_client.upload import BaseUploader
+from engine.clients.vectorsets.config import (
+    REDIS_AUTH,
+    REDIS_CLUSTER,
+    REDIS_PORT,
+    REDIS_USER,
+)
+from engine.clients.redis.helper import convert_to_redis_coords
+
+
+class RedisVsetUploader(BaseUploader):
+    client = None
+    upload_params = {}
+
+    @classmethod
+    def init_client(cls, host, distance, connection_params, upload_params):
+        redis_constructor = RedisCluster if REDIS_CLUSTER else Redis
+        cls.client = redis_constructor(
+            host=host, port=REDIS_PORT, password=REDIS_AUTH, username=REDIS_USER
+        )
+        cls.upload_params = upload_params
+
+    @classmethod
+    def upload_batch(
+        cls, ids: List[int], vectors: List[list], metadata: Optional[List[dict]]
+    ):
+        upload_params = cls.upload_params
+        hnsw_params = upload_params.get("hnsw_config")
+        M = hnsw_params.get("M", 16)
+        efc = hnsw_params.get("EF_CONSTRUCTION", 200)
+        quant = hnsw_params.get("quant")
+
+        p = cls.client.pipeline(transaction=False)
+        for i in range(len(ids)):
+            idx = ids[i]
+            vec = vectors[i]
+            vec = np.array(vec).astype(np.float32).tobytes()
+            p.execute_command("VADD", "idx", "FP32", vec, idx, quant, "M", M, "EF", efc, "CAS")
+        p.execute()
+
+    @classmethod
+    def post_upload(cls, _distance):
+        return {}