Merge remote-tracking branch 'origin/master' into update.redisearch

Author: filipecosta90

.github/workflows/continuous-benchmark.yaml

@@ -21,4 +21,11 @@ jobs:
           export GCS_SECRET=${{ secrets.GCS_SECRET }}
           export POSTGRES_PASSWORD=${{ secrets.POSTGRES_PASSWORD }}
           export POSTGRES_HOST=${{ secrets.POSTGRES_HOST }}
+
+          # Benchmark the dev branch:
+          export QDRANT_VERSION=ghcr/dev
+          bash -x tools/run_ci.sh
+
+          # Benchmark the master branch:
+          export QDRANT_VERSION=docker/master
           bash -x tools/run_ci.sh

.github/workflows/manual-benchmark.yaml

@@ -4,16 +4,16 @@ on:
   repository_dispatch:
   workflow_dispatch:
     inputs:
-        qdrant_branch:
-            description: "Branch of qdrant to benchmark"
-            default: dev
+        qdrant_version:
+            description: "Version of qdrant to benchmark (tags/v1.6.1, <commit-id>, my-branch, docker/v1.5.1, ghcr/dev)"
+            default: ghcr/dev
         dataset:
             description: "Dataset to benchmark"
             default: laion-small-clip
 
 jobs:
   runManualBenchmark:
-    name: manual benchmark - ${{ inputs.qdrant_branch }} - ${{ inputs.dataset }}
+    name: manual benchmark - ${{ inputs.qdrant_version }} - ${{ inputs.dataset }}
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
@@ -27,7 +27,7 @@ jobs:
           export GCS_SECRET=${{ secrets.GCS_SECRET }}
           export POSTGRES_PASSWORD=${{ secrets.POSTGRES_PASSWORD }}
           export POSTGRES_HOST=${{ secrets.POSTGRES_HOST }}
-          export QDRANT_VERSION=${{ inputs.qdrant_branch }}
+          export QDRANT_VERSION=${{ inputs.qdrant_version }}
           export DATASETS=${{ inputs.dataset }}
           export POSTGRES_TABLE=benchmark_manual
           bash -x tools/run_ci.sh

engine/base_client/client.py

@@ -129,3 +129,10 @@ def run_experiment(
                     print(f"\tskipping ef runtime: {ef}; #clients {parallel}")
         print("Experiment stage: Done")
         print("Results saved to: ", RESULTS_DIR)
+
+    def delete_client(self):
+        self.uploader.delete_client()
+        self.configurator.delete_client()
+
+        for s in self.searchers:
+            s.delete_client()

engine/base_client/configure.py

@@ -23,3 +23,6 @@ def configure(self, dataset: Dataset) -> Optional[dict]:
 
     def execution_params(self, distance, vector_size) -> dict:
         return {}
+
+    def delete_client(self):
+        pass

engine/clients/client_factory.py

@@ -18,6 +18,11 @@
     OpenSearchSearcher,
     OpenSearchUploader,
 )
+from engine.clients.pgvector import (
+    PgVectorConfigurator,
+    PgVectorSearcher,
+    PgVectorUploader,
+)
 from engine.clients.qdrant import QdrantConfigurator, QdrantSearcher, QdrantUploader
 from engine.clients.redis import RedisConfigurator, RedisSearcher, RedisUploader
 from engine.clients.weaviate import (
@@ -33,6 +38,7 @@
     "elasticsearch": ElasticConfigurator,
     "opensearch": OpenSearchConfigurator,
     "redis": RedisConfigurator,
+    "pgvector": PgVectorConfigurator,
 }
 
 ENGINE_UPLOADERS = {
@@ -42,6 +48,7 @@
     "elasticsearch": ElasticUploader,
     "opensearch": OpenSearchUploader,
     "redis": RedisUploader,
+    "pgvector": PgVectorUploader,
 }
 
 ENGINE_SEARCHERS = {
@@ -51,6 +58,7 @@
     "elasticsearch": ElasticSearcher,
     "opensearch": OpenSearchSearcher,
     "redis": RedisSearcher,
+    "pgvector": PgVectorSearcher,
 }
 
 

engine/clients/pgvector/__init__.py

@@ -0,0 +1,9 @@
+from engine.clients.pgvector.configure import PgVectorConfigurator
+from engine.clients.pgvector.search import PgVectorSearcher
+from engine.clients.pgvector.upload import PgVectorUploader
+
+__all__ = [
+    "PgVectorConfigurator",
+    "PgVectorSearcher",
+    "PgVectorUploader",
+]

engine/clients/pgvector/config.py

@@ -0,0 +1,15 @@
+PGVECTOR_PORT = 9200
+PGVECTOR_DB = "postgres"
+PGVECTOR_USER = "postgres"
+PGVECTOR_PASSWORD = "passwd"
+
+
+def get_db_config(host, connection_params):
+    return {
+        "host": host or "localhost",
+        "dbname": PGVECTOR_DB,
+        "user": PGVECTOR_USER,
+        "password": PGVECTOR_PASSWORD,
+        "autocommit": True,
+        **connection_params,
+    }

engine/clients/pgvector/configure.py

@@ -0,0 +1,55 @@
+import pgvector.psycopg
+import psycopg
+
+from benchmark.dataset import Dataset
+from engine.base_client import IncompatibilityError
+from engine.base_client.configure import BaseConfigurator
+from engine.base_client.distances import Distance
+from engine.clients.pgvector.config import get_db_config
+
+
+class PgVectorConfigurator(BaseConfigurator):
+    DISTANCE_MAPPING = {
+        Distance.L2: "vector_l2_ops",
+        Distance.COSINE: "vector_cosine_ops",
+    }
+
+    def __init__(self, host, collection_params: dict, connection_params: dict):
+        super().__init__(host, collection_params, connection_params)
+        self.conn = psycopg.connect(**get_db_config(host, connection_params))
+        print("configure connection created")
+        self.conn.execute("CREATE EXTENSION IF NOT EXISTS vector;")
+        pgvector.psycopg.register_vector(self.conn)
+
+    def clean(self):
+        self.conn.execute(
+            "DROP TABLE IF EXISTS items CASCADE;",
+        )
+
+    def recreate(self, dataset: Dataset, collection_params):
+        if dataset.config.distance == Distance.DOT:
+            raise IncompatibilityError
+
+        self.conn.execute(
+            f"""CREATE TABLE items (
+                id SERIAL PRIMARY KEY,
+                embedding vector({dataset.config.vector_size}) NOT NULL
+            );"""
+        )
+        self.conn.execute("ALTER TABLE items ALTER COLUMN embedding SET STORAGE PLAIN")
+
+        try:
+            hnsw_distance_type = self.DISTANCE_MAPPING[dataset.config.distance]
+        except KeyError:
+            raise IncompatibilityError(
+                f"Unsupported distance metric: {dataset.config.distance}"
+            )
+
+        self.conn.execute(
+            f"CREATE INDEX on items USING hnsw(embedding {hnsw_distance_type}) WITH (m = {collection_params['hnsw_config']['m']}, ef_construction = {collection_params['hnsw_config']['ef_construct']})"
+        )
+
+        self.conn.close()
+
+    def delete_client(self):
+        self.conn.close()

engine/clients/pgvector/parser.py

@@ -0,0 +1,46 @@
+import json
+from typing import Any, List, Optional
+
+from engine.base_client import IncompatibilityError
+from engine.base_client.parser import BaseConditionParser, FieldValue
+
+
+class PgVectorConditionParser(BaseConditionParser):
+    def build_condition(
+        self, and_subfilters: Optional[List[Any]], or_subfilters: Optional[List[Any]]
+    ) -> Optional[Any]:
+        clauses = []
+        if or_subfilters is not None and len(or_subfilters) > 0:
+            clauses.append(f"( {' OR '.join(or_subfilters)} )")
+        if and_subfilters is not None and len(and_subfilters) > 0:
+            clauses.append(f"( {' AND '.join(or_subfilters)} )")
+
+        return " AND ".join(clauses)
+
+    def build_exact_match_filter(self, field_name: str, value: FieldValue) -> Any:
+        raise f"{field_name} == {json.dumps(value)}"
+
+    def build_range_filter(
+        self,
+        field_name: str,
+        lt: Optional[FieldValue],
+        gt: Optional[FieldValue],
+        lte: Optional[FieldValue],
+        gte: Optional[FieldValue],
+    ) -> Any:
+        clauses = []
+        if lt is not None:
+            clauses.append(f"{field_name} < {lt}")
+        if gt is not None:
+            clauses.append(f"{field_name} > {gt}")
+        if lte is not None:
+            clauses.append(f"{field_name} <= {lte}")
+        if gte is not None:
+            clauses.append(f"{field_name} >= {gte}")
+        return f"( {' AND '.join(clauses)} )"
+
+    def build_geo_filter(
+        self, field_name: str, lat: float, lon: float, radius: float
+    ) -> Any:
+        # TODO: Implement this
+        raise IncompatibilityError

engine/clients/pgvector/search.py

@@ -0,0 +1,50 @@
+import multiprocessing as mp
+from typing import List, Tuple
+
+import numpy as np
+import psycopg
+from pgvector.psycopg import register_vector
+
+from engine.base_client.distances import Distance
+from engine.base_client.search import BaseSearcher
+from engine.clients.pgvector.config import get_db_config
+from engine.clients.pgvector.parser import PgVectorConditionParser
+
+
+class PgVectorSearcher(BaseSearcher):
+    conn = None
+    cur = None
+    distance = None
+    search_params = {}
+    parser = PgVectorConditionParser()
+
+    @classmethod
+    def init_client(cls, host, distance, connection_params: dict, search_params: dict):
+        cls.conn = psycopg.connect(**get_db_config(host, connection_params))
+        register_vector(cls.conn)
+        cls.cur = cls.conn.cursor()
+        cls.distance = distance
+        cls.search_params = search_params["search_params"]
+
+    @classmethod
+    def search_one(cls, vector, meta_conditions, top) -> List[Tuple[int, float]]:
+        cls.cur.execute(f"SET hnsw.ef_search = {cls.search_params['hnsw_ef']}")
+
+        if cls.distance == Distance.COSINE:
+            query = f"SELECT id, embedding <=> %s AS _score FROM items ORDER BY _score LIMIT {top};"
+        elif cls.distance == Distance.L2:
+            query = f"SELECT id, embedding <-> %s AS _score FROM items ORDER BY _score LIMIT {top};"
+        else:
+            raise NotImplementedError(f"Unsupported distance metric {cls.distance}")
+
+        cls.cur.execute(
+            query,
+            (np.array(vector),),
+        )
+        return cls.cur.fetchall()
+
+    @classmethod
+    def delete_client(cls):
+        if cls.cur:
+            cls.cur.close()
+            cls.conn.close()