vectordb-bench 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

vectordb_bench/backend/clients/pgvector/config.py

@@ -0,0 +1,49 @@
+from pydantic import BaseModel, SecretStr
+from ..api import DBConfig, DBCaseConfig, MetricType
+
+POSTGRE_URL_PLACEHOLDER = "postgresql://%s:%s@%s/%s"
+
+class PgVectorConfig(DBConfig):
+    user_name: SecretStr = "postgres"
+    password: SecretStr
+    url: SecretStr
+    db_name: str
+
+    def to_dict(self) -> dict:
+        user_str = self.user_name.get_secret_value()
+        pwd_str = self.password.get_secret_value()
+        url_str = self.url.get_secret_value()
+        return {
+            "url" : POSTGRE_URL_PLACEHOLDER%(user_str, pwd_str, url_str, self.db_name)
+        }
+
+class PgVectorIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    lists: int | None = 1000
+    probes: int | None = 10
+
+    def parse_metric(self) -> str: 
+        if self.metric_type == MetricType.L2:
+            return "vector_l2_ops"
+        elif self.metric_type == MetricType.IP:
+            return "vector_ip_ops"
+        return "vector_cosine_ops"
+    
+    def parse_metric_fun_str(self) -> str: 
+        if self.metric_type == MetricType.L2:
+            return "l2_distance"
+        elif self.metric_type == MetricType.IP:
+            return "max_inner_product"
+        return "cosine_distance"
+
+    def index_param(self) -> dict:
+        return {
+            "lists" : self.lists,
+            "metric" : self.parse_metric()
+        }
+    
+    def search_param(self) -> dict:
+        return {
+            "probes" : self.probes,
+            "metric_fun" : self.parse_metric_fun_str()
+        }

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -0,0 +1,171 @@
+"""Wrapper around the Pgvector vector database over VectorDB"""
+
+import logging
+import time
+from contextlib import contextmanager
+from typing import Any, Type
+from functools import wraps
+
+from ..api import VectorDB, DBConfig, DBCaseConfig, IndexType
+from pgvector.sqlalchemy import Vector
+from .config import PgVectorConfig, PgVectorIndexConfig
+from sqlalchemy import (
+    MetaData,
+    create_engine,
+    insert,
+    select,
+    Index,
+    Table,
+    text,
+    Column,
+    Float, 
+    Integer
+)
+from sqlalchemy.orm import (
+    declarative_base, 
+    mapped_column, 
+    Session
+)
+
+log = logging.getLogger(__name__) 
+
+class PgVector(VectorDB):
+    """ Use SQLAlchemy instructions"""
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: DBCaseConfig,
+        collection_name: str = "PgVectorCollection",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+
+        self._index_name = "pqvector_index"
+        self._primary_field = "id"
+        self._vector_field = "embedding"
+
+        # construct basic units
+        pg_engine = create_engine(**self.db_config)
+        Base = declarative_base()
+        pq_metadata = Base.metadata
+        pq_metadata.reflect(pg_engine) 
+        
+        # create vector extension
+        with pg_engine.connect() as conn: 
+            conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
+            conn.commit()
+        
+        self.pg_table = self._get_table_schema(pq_metadata)
+        if drop_old and self.table_name in pq_metadata.tables:
+            log.info(f"Pgvector client drop table : {self.table_name}")
+            # self.pg_table.drop(pg_engine, checkfirst=True)
+            pq_metadata.drop_all(pg_engine)
+            self._create_table(dim, pg_engine)
+
+    
+    @classmethod
+    def config_cls(cls) -> Type[DBConfig]:
+        return PgVectorConfig
+
+    @classmethod
+    def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+        return PgVectorIndexConfig
+
+    @contextmanager
+    def init(self) -> None:
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+        self.pg_engine = create_engine(**self.db_config)
+
+        Base = declarative_base()
+        pq_metadata = Base.metadata
+        pq_metadata.reflect(self.pg_engine) 
+        self.pg_session = Session(self.pg_engine)
+        self.pg_table = self._get_table_schema(pq_metadata)
+        yield 
+        self.pg_session = None
+        self.pg_engine = None 
+        del (self.pg_session)
+        del (self.pg_engine)
+    
+    def ready_to_load(self):
+        pass
+
+    def optimize(self):
+        pass
+
+    def ready_to_search(self):
+        pass
+
+    def _get_table_schema(self, pq_metadata):
+        return Table(
+            self.table_name,
+            pq_metadata,
+            Column(self._primary_field, Integer, primary_key=True),
+            Column(self._vector_field, Vector(self.dim)),
+            extend_existing=True
+        )
+    
+    def _create_index(self, pg_engine):
+        index_param = self.case_config.index_param()
+        index = Index(self._index_name, self.pg_table.c.embedding,
+            postgresql_using='ivfflat',
+            postgresql_with={'lists': index_param["lists"]},
+            postgresql_ops={'embedding': index_param["metric"]}
+        )
+        index.drop(pg_engine, checkfirst = True)
+        index.create(pg_engine)
+
+    def _create_table(self, dim, pg_engine : int):
+        try:
+            # create table
+            self.pg_table.create(bind = pg_engine, checkfirst = True)
+            # create vec index
+            self._create_index(pg_engine)
+        except Exception as e:
+            log.warning(f"Failed to create pgvector table: {self.table_name} error: {e}")
+            raise e from None
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> (int, Exception):
+        try:
+            items = [dict(id = metadata[i], embedding=embeddings[i]) for i in range(len(metadata))]
+            self.pg_session.execute(insert(self.pg_table), items)
+            self.pg_session.commit()
+            return len(metadata), None
+        except Exception as e:
+            log.warning(f"Failed to insert data into pgvector table ({self.table_name}), error: {e}")   
+            return 0, e
+
+    def search_embedding(        
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+    ) -> list[int]:
+        assert self.pg_table is not None
+        search_param =self.case_config.search_param()
+        with self.pg_engine.connect() as conn: 
+            conn.execute(text(f'SET ivfflat.probes = {search_param["probes"]}'))
+            conn.commit()
+        op_fun = getattr(self.pg_table.c.embedding, search_param["metric_fun"])
+        if filters:
+            res = self.pg_session.scalars(select(self.pg_table).order_by(op_fun(query)).filter(self.pg_table.c.id > filters.get('id')).limit(k))
+        else: 
+            res = self.pg_session.scalars(select(self.pg_table).order_by(op_fun(query)).limit(k))
+        return list(res)
+        

vectordb_bench/backend/clients/pinecone/config.py

@@ -2,9 +2,9 @@ from pydantic import BaseModel, SecretStr
 from ..api import DBConfig
 
 
-class PineconeConfig(DBConfig, BaseModel):
-    api_key: SecretStr | None = None
-    environment: SecretStr | None = None
+class PineconeConfig(DBConfig):
+    api_key: SecretStr
+    environment: SecretStr
     index_name: str
 
     def to_dict(self) -> dict:

vectordb_bench/backend/clients/pinecone/pinecone.py

@@ -2,7 +2,7 @@
 
 import logging
 from contextlib import contextmanager
-from typing import Any, Type
+from typing import Type
 
 from ..api import VectorDB, DBConfig, DBCaseConfig, EmptyDBCaseConfig, IndexType
 from .config import PineconeConfig
@@ -20,6 +20,7 @@ class Pinecone(VectorDB):
         db_config: dict,
         db_case_config: DBCaseConfig,
         drop_old: bool = False,
+        **kwargs,
     ):
         """Initialize wrapper around the milvus vector database."""
         self.index_name = db_config["index_name"]
@@ -69,24 +70,30 @@ class Pinecone(VectorDB):
     def ready_to_load(self):
         pass
 
-    def ready_to_search(self):
+    def optimize(self):
         pass
 
     def insert_embeddings(
         self,
         embeddings: list[list[float]],
         metadata: list[int],
-    ) -> list[str]:
+        **kwargs,
+    ) -> (int, Exception):
         assert len(embeddings) == len(metadata)
-        for batch_start_offset in range(0, len(embeddings), self.batch_size):
-            batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
-            insert_datas = []
-            for i in range(batch_start_offset, batch_end_offset):
-                insert_data = (str(metadata[i]), embeddings[i], {
-                            self._metadata_key: metadata[i]})
-                insert_datas.append(insert_data)
-            self.index.upsert(insert_datas)
-        return len(embeddings)
+        insert_count = 0
+        try:
+            for batch_start_offset in range(0, len(embeddings), self.batch_size):
+                batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
+                insert_datas = []
+                for i in range(batch_start_offset, batch_end_offset):
+                    insert_data = (str(metadata[i]), embeddings[i], {
+                                self._metadata_key: metadata[i]})
+                    insert_datas.append(insert_data)
+                self.index.upsert(insert_datas)
+                insert_count += batch_end_offset - batch_start_offset
+        except Exception as e:
+            return (insert_count, e)
+        return (len(embeddings), None)
 
     def search_embedding(
         self,
@@ -94,7 +101,6 @@ class Pinecone(VectorDB):
         k: int = 100,
         filters: dict | None = None,
         timeout: int | None = None,
-        **kwargs: Any,
     ) -> list[tuple[int, float]]:
         if filters is None:
             pinecone_filters = {}

vectordb_bench/backend/clients/qdrant_cloud/config.py

@@ -1,16 +1,33 @@
 from pydantic import BaseModel, SecretStr
 
-from ..api import DBConfig
+from ..api import DBConfig, DBCaseConfig, MetricType
+from qdrant_client.models import Distance
 
 
-class QdrantConfig(DBConfig, BaseModel):
-    url: SecretStr | None = None
-    api_key: SecretStr | None = None
-    prefer_grpc: bool = True
+class QdrantConfig(DBConfig):
+    url: SecretStr
+    api_key: SecretStr
 
     def to_dict(self) -> dict:
         return {
             "url": self.url.get_secret_value(),
             "api_key": self.api_key.get_secret_value(),
-            "prefer_grpc": self.prefer_grpc,
+            "prefer_grpc": True,
         }
+
+class QdrantIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return Distance.EUCLID
+        elif self.metric_type == MetricType.IP:
+            return Distance.DOT
+        return Distance.COSINE
+
+    def index_param(self) -> dict:
+        params = {"distance": self.parse_metric()}
+        return params
+
+    def search_param(self) -> dict:
+        return {}

vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py

@@ -3,13 +3,12 @@
 import logging
 import time
 from contextlib import contextmanager
-from typing import Any, Type
+from typing import Type
 
-from ..api import VectorDB, DBConfig, DBCaseConfig, EmptyDBCaseConfig, IndexType
-from .config import QdrantConfig
+from ..api import VectorDB, DBConfig, DBCaseConfig, IndexType
+from .config import QdrantConfig, QdrantIndexConfig
 from qdrant_client.http.models import (
     CollectionStatus,
-    Distance,
     VectorParams,
     PayloadSchemaType,
     Batch,
@@ -32,6 +31,7 @@ class QdrantCloud(VectorDB):
         db_case_config: DBCaseConfig,
         collection_name: str = "QdrantCloudCollection",
         drop_old: bool = False,
+        **kwargs,
     ):
         """Initialize wrapper around the QdrantCloud vector database."""
         self.db_config = db_config
@@ -55,7 +55,7 @@ class QdrantCloud(VectorDB):
 
     @classmethod
     def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
-        return EmptyDBCaseConfig
+        return QdrantIndexConfig
 
     @contextmanager
     def init(self) -> None:
@@ -74,7 +74,7 @@ class QdrantCloud(VectorDB):
         pass
 
 
-    def ready_to_search(self):
+    def optimize(self):
         assert self.qdrant_client, "Please call self.init() before"
         # wait for vectors to be fully indexed
         SECONDS_WAITING_FOR_INDEXING_API_CALL = 5
@@ -97,7 +97,7 @@ class QdrantCloud(VectorDB):
         try:
             qdrant_client.create_collection(
                 collection_name=self.collection_name,
-                vectors_config=VectorParams(size=dim, distance=Distance.EUCLID)
+                vectors_config=VectorParams(size=dim, distance=self.case_config.index_param()["distance"])
             )
 
             qdrant_client.create_payload_index(
@@ -116,8 +116,8 @@ class QdrantCloud(VectorDB):
         self,
         embeddings: list[list[float]],
         metadata: list[int],
-        **kwargs: Any,
-    ) -> list[str]:
+        **kwargs,
+    ) -> (int, Exception):
         """Insert embeddings into Milvus. should call self.init() first"""
         assert self.qdrant_client is not None
         try:
@@ -127,11 +127,11 @@ class QdrantCloud(VectorDB):
                 wait=True,
                 points=Batch(ids=metadata, payloads=[{self._primary_field: v} for v in metadata], vectors=embeddings)
             )
-
-            return len(metadata)
         except Exception as e:
             log.info(f"Failed to insert data, {e}")
-            raise e from None
+            return 0, e
+        else:
+            return len(metadata), None
 
     def search_embedding(
         self,
@@ -139,7 +139,6 @@ class QdrantCloud(VectorDB):
         k: int = 100,
         filters: dict | None = None,
         timeout: int | None = None,
-        **kwargs: Any,
     ) -> list[int]:
         """Perform a search on a query embedding and return results with score.
         Should call self.init() first.

vectordb_bench/backend/clients/weaviate_cloud/config.py

@@ -4,9 +4,9 @@ import weaviate
 from ..api import DBConfig, DBCaseConfig, MetricType
 
 
-class WeaviateConfig(DBConfig, BaseModel):
-    url: SecretStr | None = None
-    api_key: SecretStr | None = None
+class WeaviateConfig(DBConfig):
+    url: SecretStr
+    api_key: SecretStr
 
     def to_dict(self) -> dict:
         return {

vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py

@@ -1,7 +1,7 @@
 """Wrapper around the Weaviate vector database over VectorDB"""
 
 import logging
-from typing import Any, Iterable, Type
+from typing import Iterable, Type
 from contextlib import contextmanager
 
 from weaviate.exceptions import WeaviateBaseError
@@ -21,6 +21,7 @@ class WeaviateCloud(VectorDB):
         db_case_config: DBCaseConfig,
         collection_name: str = "VectorDBBenchCollection",
         drop_old: bool = False,
+        **kwargs,
     ):
         """Initialize wrapper around the weaviate vector database."""
         self.db_config = db_config
@@ -70,7 +71,7 @@ class WeaviateCloud(VectorDB):
         """Should call insert first, do nothing"""
         pass
 
-    def ready_to_search(self):
+    def optimize(self):
         assert self.client.schema.exists(self.collection_name)
         self.client.schema.update_config(self.collection_name, {"vectorIndexConfig": self.case_config.search_param() } )
 
@@ -98,11 +99,11 @@ class WeaviateCloud(VectorDB):
         self,
         embeddings: Iterable[list[float]],
         metadata: list[int],
-        **kwargs: Any,
-    ) -> int:
+        **kwargs,
+    ) -> (int, Exception):
         """Insert embeddings into Weaviate"""
         assert self.client.schema.exists(self.collection_name)
-
+        insert_count = 0
         try:
             with self.client.batch as batch:
                 batch.batch_size = len(metadata)
@@ -114,10 +115,11 @@ class WeaviateCloud(VectorDB):
                         class_name=self.collection_name,
                         vector=embeddings[i]
                     ))
-                return len(res)
+                    insert_count += 1
+                return (len(res), None)
         except WeaviateBaseError as e:
             log.warning(f"Failed to insert data, error: {str(e)}")
-            raise e from None
+            return (insert_count, e)
 
     def search_embedding(
         self,
@@ -125,7 +127,6 @@ class WeaviateCloud(VectorDB):
         k: int = 100,
         filters: dict | None = None,
         timeout: int | None = None,
-        **kwargs: Any,
     ) -> list[int]:
         """Perform a search on a query embedding and return results with distance.
         Should call self.init() first.

vectordb_bench/backend/clients/zilliz_cloud/config.py

@@ -1,12 +1,13 @@
-from pydantic import BaseModel, SecretStr
+from pydantic import SecretStr
+
 from ..api import DBCaseConfig, DBConfig
 from ..milvus.config import MilvusIndexConfig, IndexType
 
 
-class ZillizCloudConfig(DBConfig, BaseModel):
-    uri: SecretStr | None = None
+class ZillizCloudConfig(DBConfig):
+    uri: SecretStr
     user: str
-    password: SecretStr | None = None
+    password: SecretStr
 
     def to_dict(self) -> dict:
         return {

vectordb_bench/backend/clients/zilliz_cloud/zilliz_cloud.py

@@ -14,7 +14,8 @@ class ZillizCloud(Milvus):
         db_case_config: DBCaseConfig,
         collection_name: str = "ZillizCloudVectorDBBench",
         drop_old: bool = False,
-        name: str = "ZillizCloud"
+        name: str = "ZillizCloud",
+        **kwargs, 
     ):
         super().__init__(
             dim=dim,
@@ -23,6 +24,7 @@ class ZillizCloud(Milvus):
             collection_name=collection_name,
             drop_old=drop_old,
             name=name,
+            **kwargs,
         )
 
     @classmethod