elastic

Author: generall

datasets/random-100/neighbours.jsonl

@@ -0,0 +1,10 @@
+[0]
+[1]
+[2]
+[3]
+[4]
+[5]
+[6]
+[7]
+[8]
+[9]

engine/clients/client_factory.py

@@ -8,7 +8,10 @@
     BaseUploader,
 )
 from engine.clients.milvus import MilvusConfigurator, MilvusSearcher, MilvusUploader
-from engine.clients.qdrant import QdrantConfigurator, QdrantSearcher, QdrantUploader
+from engine.clients.elasticsearch.configure import ElasticConfigurator
+from engine.clients.elasticsearch.search import ElasticSearcher
+from engine.clients.elasticsearch.upload import ElasticUploader
+from engine.clients.qdrant import QdrantConfigurator, QdrantUploader, QdrantSearcher
 from engine.clients.weaviate import (
     WeaviateConfigurator,
     WeaviateSearcher,
@@ -19,18 +22,21 @@
     "qdrant": QdrantConfigurator,
     "weaviate": WeaviateConfigurator,
     "milvus": MilvusConfigurator,
+    "elastic": ElasticConfigurator,
 }
 
 ENGINE_UPLOADERS = {
     "qdrant": QdrantUploader,
     "weaviate": WeaviateUploader,
     "milvus": MilvusUploader,
+    "elastic": ElasticUploader,
 }
 
 ENGINE_SEARCHERS = {
     "qdrant": QdrantSearcher,
     "weaviate": WeaviateSearcher,
     "milvus": MilvusSearcher,
+    "elastic": ElasticSearcher,
 }
 
 

engine/elasticsearch-8.3.1/client/config.py renamed to engine/clients/elasticsearch/__init__.py

@@ -1,6 +1,3 @@
-import os
-
-ELASTIC_HOST = os.environ.get("SERVER_HOST", "elastic_server")
 ELASTIC_PORT = 9200
 ELASTIC_INDEX = "bench"
 ELASTIC_USER = "elastic"

engine/clients/elasticsearch/configure.py

@@ -0,0 +1,60 @@
+from elasticsearch import Elasticsearch, NotFoundError
+
+from engine.base_client.configure import BaseConfigurator
+from engine.base_client.distances import Distance
+from engine.clients.elasticsearch import ELASTIC_USER, ELASTIC_PASSWORD, ELASTIC_PORT, ELASTIC_INDEX
+
+
+class ElasticConfigurator(BaseConfigurator):
+    DISTANCE_MAPPING = {
+        Distance.L2_SQUARED: "l2_norm",
+        Distance.COSINE: "cosine",
+        Distance.DOT: "dot_product",
+    }
+
+    def __init__(self, host, collection_params: dict, connection_params: dict):
+        super().__init__(host, collection_params, connection_params)
+
+        self.client = Elasticsearch(
+            f"http://{host}:{ELASTIC_PORT}",
+            basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),
+            **{
+                **{
+                    "verify_certs": False,
+                    "request_timeout": 90,
+                    "retry_on_timeout": True,
+                },
+                **connection_params
+            }
+        )
+
+    def clean(self):
+        try:
+            self.client.indices.delete(index=ELASTIC_INDEX)
+        except NotFoundError:
+            pass
+
+    def recreate(
+            self, distance, vector_size, collection_params,
+    ):
+        self.client.indices.create(
+            index=ELASTIC_INDEX,
+            mappings={
+                "properties": {
+                    "vector": {
+                        "type": "dense_vector",
+                        "dims": vector_size,
+                        "index": True,
+                        "similarity": self.DISTANCE_MAPPING[distance],
+                        "index_options": {
+                            **{
+                                "type": "hnsw",
+                                "m": 16,
+                                "ef_construction": 100,
+                            },
+                            **collection_params.get("index_options"),
+                        },
+                    }
+                }
+            }
+        )

engine/clients/elasticsearch/search.py

@@ -0,0 +1,44 @@
+import uuid
+from typing import Tuple, List
+
+from elasticsearch import Elasticsearch
+
+from engine.base_client.search import BaseSearcher
+from engine.clients.elasticsearch import ELASTIC_PORT, ELASTIC_USER, ELASTIC_PASSWORD, ELASTIC_INDEX
+
+
+class ElasticSearcher(BaseSearcher):
+    search_params = {}
+    client: Elasticsearch = None
+
+    @classmethod
+    def init_client(cls, host, connection_params: dict, search_params: dict):
+        cls.client: Elasticsearch = Elasticsearch(
+            f"http://{host}:{ELASTIC_PORT}",
+            basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),
+            **{
+                **{
+                    "verify_certs": False,
+                    "request_timeout": 90,
+                    "retry_on_timeout": True,
+                },
+                **connection_params
+            }
+        )
+        cls.search_params = search_params
+
+    @classmethod
+    def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
+        res = cls.client.knn_search(
+            index=ELASTIC_INDEX,
+            knn={
+                "field": "vector",
+                "query_vector": vector,
+                "k": 10,
+                **{
+                    "num_candidates": 100,
+                    **cls.search_params
+                }
+            },
+        )
+        return [(uuid.UUID(hex=hit["_id"]).int, hit["_score"]) for hit in res["hits"]["hits"]]

engine/clients/elasticsearch/upload.py

@@ -0,0 +1,52 @@
+import uuid
+from typing import Optional, List
+
+from elasticsearch import Elasticsearch
+
+from engine.base_client.upload import BaseUploader
+from engine.clients.elasticsearch import ELASTIC_PORT, ELASTIC_USER, ELASTIC_PASSWORD, ELASTIC_INDEX
+
+
+class ElasticUploader(BaseUploader):
+    client = None
+    upload_params = {}
+
+    @classmethod
+    def init_client(cls, host, connection_params, upload_params):
+        cls.client: Elasticsearch = Elasticsearch(
+            f"http://{host}:{ELASTIC_PORT}",
+            basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),
+            **{
+                **{
+                    "verify_certs": False,
+                    "request_timeout": 90,
+                    "retry_on_timeout": True,
+                },
+                **connection_params
+            }
+        )
+        cls.upload_params = upload_params
+
+    @classmethod
+    def upload_batch(
+        cls, ids: List[int], vectors: List[list], metadata: Optional[List[dict]]
+    ):
+        operations = []
+        for idx, vector in zip(ids, vectors):
+            vector_id = uuid.UUID(int=idx).hex
+
+            operations.append({"index": {"_id": vector_id}})
+            operations.append({"vector": vector})
+
+        cls.client.bulk(
+            index=ELASTIC_INDEX,
+            operations=operations,
+        )
+
+    @classmethod
+    def post_upload(cls):
+        cls.client.indices.forcemerge(
+            index=ELASTIC_INDEX,
+            wait_for_completion=True
+        )
+        return {}

engine/elasticsearch-8.3.1/client.Dockerfile

@@ -2,23 +2,13 @@ version: '3.5'
 
 services:
   es:
-    image: docker.elastic.co/elasticsearch/elasticsearch:8.3.1
-    volumes:
-      - es_data:/usr/share/elasticsearch/data
+    image: docker.elastic.co/elasticsearch/elasticsearch:8.3.3
     environment:
       ELASTIC_PASSWORD: "passwd"
       KIBANA_PASSWORD: "passwd"
       SERVER_SSL_ENABLED: "false"
       discovery.type: "single-node"
+      xpack.security.enabled: "false"
     ports:
       - "9200:9200"
       - "9300:9300"
-    networks:
-      - elastic
-
-networks:
-  elastic:
-    driver: bridge
-
-volumes:
-  es_data: