add redis

Author: generall

benchmark/config_read.py

@@ -0,0 +1,28 @@
+import glob
+import json
+import os
+
+from benchmark import ROOT_DIR, DATASETS_DIR
+
+
+def read_engine_configs() -> dict:
+    all_configs = {}
+    engines_config_dir = os.path.join(ROOT_DIR, "experiments", "configurations")
+    config_files = glob.glob(os.path.join(engines_config_dir, "*.json"))
+    for config_file in config_files:
+        with open(config_file, "r") as fd:
+            configs = json.load(fd)
+            for config in configs:
+                all_configs[config["name"]] = config
+
+    return all_configs
+
+
+def read_dataset_config():
+    all_configs = {}
+    datasets_config_path = DATASETS_DIR / "datasets.json"
+    with open(datasets_config_path, "r") as fd:
+        configs = json.load(fd)
+        for config in configs:
+            all_configs[config["name"]] = config
+    return all_configs

engine/base_client/search.py

@@ -85,9 +85,9 @@ def search_all(
         total_time = time.perf_counter() - start
         return {
             "total_time": total_time,
-            "mean_time": np.std(latencies),
+            "mean_time": np.mean(latencies),
             "mean_precisions": np.mean(precisions),
-            "std_time": np.mean(latencies),
+            "std_time": np.std(latencies),
             "min_time": np.min(latencies),
             "max_time": np.max(latencies),
             "rps": len(latencies) / total_time,

engine/clients/client_factory.py

@@ -12,6 +12,9 @@
 from engine.clients.elasticsearch.upload import ElasticUploader
 from engine.clients.milvus import MilvusConfigurator, MilvusSearcher, MilvusUploader
 from engine.clients.qdrant import QdrantConfigurator, QdrantSearcher, QdrantUploader
+from engine.clients.redis_single_node.configure import RedisConfigurator
+from engine.clients.redis_single_node.search import RedisSearcher
+from engine.clients.redis_single_node.upload import RedisUploader
 from engine.clients.weaviate import (
     WeaviateConfigurator,
     WeaviateSearcher,
@@ -23,20 +26,23 @@
     "weaviate": WeaviateConfigurator,
     "milvus": MilvusConfigurator,
     "elastic": ElasticConfigurator,
+    "redis": RedisConfigurator,
 }
 
 ENGINE_UPLOADERS = {
     "qdrant": QdrantUploader,
     "weaviate": WeaviateUploader,
     "milvus": MilvusUploader,
     "elastic": ElasticUploader,
+    "redis": RedisUploader,
 }
 
 ENGINE_SEARCHERS = {
     "qdrant": QdrantSearcher,
     "weaviate": WeaviateSearcher,
     "milvus": MilvusSearcher,
     "elastic": ElasticSearcher,
+    "redis": RedisSearcher,
 }
 
 

engine/clients/qdrant/upload.py

@@ -32,18 +32,21 @@ def upload_batch(
 
     @classmethod
     def post_upload(cls, _distance):
-        cls.wait_collection_green()
-        time.sleep(1.0)
         cls.wait_collection_green()
         return {}
 
     @classmethod
     def wait_collection_green(cls):
-        wait_time = 1.0
+        wait_time = 5.0
         total = 0
-        collection_info = cls.client.get_collection(QDRANT_COLLECTION_NAME)
-        while collection_info.status != CollectionStatus.GREEN:
+        while True:
             time.sleep(wait_time)
             total += wait_time
             collection_info = cls.client.get_collection(QDRANT_COLLECTION_NAME)
+            if collection_info.status != CollectionStatus.GREEN:
+                continue
+            time.sleep(wait_time)
+            collection_info = cls.client.get_collection(QDRANT_COLLECTION_NAME)
+            if collection_info.status == CollectionStatus.GREEN:
+                break
         return total

engine/clients/redis_single_node/__init__.py

@@ -0,0 +1 @@
+REDIS_INDEX_NAME = "benchmark"

engine/clients/redis_single_node/config.py

@@ -0,0 +1,50 @@
+import redis
+from redis.commands.search.field import VectorField
+
+from engine.base_client.configure import BaseConfigurator
+from engine.base_client.distances import Distance
+
+
+class RedisConfigurator(BaseConfigurator):
+    DISTANCE_MAPPING = {
+        Distance.L2: "L2",
+        Distance.COSINE: "COSINE",
+        Distance.DOT: "IP",
+    }
+
+    def __init__(self, host, collection_params: dict, connection_params: dict):
+        super().__init__(host, collection_params, connection_params)
+
+        self.client = redis.Redis(host=host, port=6379, db=0)
+
+    def clean(self):
+        index = self.client.ft()
+        try:
+            index.dropindex(delete_documents=True)
+        except redis.ResponseError as e:
+            print(e)
+
+    def recreate(
+        self,
+        distance,
+        vector_size,
+        collection_params,
+    ):
+        self.clean()
+        index = self.client.ft()
+        index.create_index(
+            fields=[VectorField(
+                name="vector",
+                algorithm="HNSW",
+                attributes={
+                    "TYPE": "FLOAT32",
+                    "DIM": vector_size,
+                    "DISTANCE_METRIC": self.DISTANCE_MAPPING[distance],
+                    **self.collection_params.get('hnsw_config', {}),
+                }
+            )]
+        )
+
+
+if __name__ == '__main__':
+    pass

engine/clients/redis_single_node/configure.py

@@ -0,0 +1,34 @@
+from typing import List, Tuple
+
+import numpy as np
+import redis
+from redis.commands.search.query import Query
+
+from engine.base_client.search import BaseSearcher
+
+
+class RedisSearcher(BaseSearcher):
+    search_params = {}
+    client = None
+
+    @classmethod
+    def init_client(cls, host, distance, connection_params: dict, search_params: dict):
+        cls.client = redis.Redis(host=host, port=6379, db=0)
+        cls.search_params = search_params
+
+    @classmethod
+    def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
+        q = Query(f'*=>[KNN $K @vector $vec_param EF_RUNTIME $EF AS vector_score]') \
+            .sort_by('vector_score') \
+            .paging(0, top) \
+            .return_fields('vector_score') \
+            .dialect(2)
+        params_dict = {
+            "vec_param": np.array(vector).astype(np.float32).tobytes(),
+            "K": top,
+            "EF": cls.search_params['search_params']["ef"],
+        }
+
+        results = cls.client.ft().search(q, query_params=params_dict)
+
+        return [(int(result.id), float(result.vector_score)) for result in results.docs]

engine/clients/redis_single_node/search.py

@@ -0,0 +1,39 @@
+from typing import List, Optional
+
+import numpy as np
+import redis
+
+from engine.base_client.upload import BaseUploader
+
+
+class RedisUploader(BaseUploader):
+    client = None
+    upload_params = {}
+
+    @classmethod
+    def init_client(cls, host, distance, connection_params, upload_params):
+        cls.client = redis.Redis(host=host, port=6379, db=0)
+        cls.upload_params = upload_params
+
+    @classmethod
+    def upload_batch(
+        cls, ids: List[int], vectors: List[list], metadata: Optional[List[dict]]
+    ):
+
+        p = cls.client.pipeline(transaction=False)
+        for i in range(len(ids)):
+            idx = ids[i]
+            vec = vectors[i]
+            meta = metadata[i] if metadata else {}
+            meta = meta or {}
+            cls.client.hset(str(idx), mapping={
+                "vector": np.array(vec).astype(np.float32).tobytes(),
+                **meta
+            })
+        p.execute()
+
+    @classmethod
+    def post_upload(cls, _distance):
+        return {}
+
+

engine/clients/redis_single_node/upload.py

@@ -12,3 +12,7 @@ services:
     ports:
       - "9200:9200"
       - "9300:9300"
+    deploy:
+      resources:
+        limits:
+          memory: 25Gb