vectordb-bench 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

vectordb_bench/backend/clients/__init__.py

@@ -54,8 +54,8 @@ class DB(Enum):
             return ElasticCloud
 
         if self == DB.QdrantCloud:
-            from .qdrant_cloud.qdrant_cloud import QdrantClient
-            return QdrantClient
+            from .qdrant_cloud.qdrant_cloud import QdrantCloud
+            return QdrantCloud
 
         if self == DB.WeaviateCloud:
             from .weaviate_cloud.weaviate_cloud import WeaviateCloud
@@ -142,8 +142,8 @@ class DB(Enum):
             return WeaviateIndexConfig
 
         if self == DB.PgVector:
-            from .pgvector.config import PgVectorIndexConfig
-            return PgVectorIndexConfig
+            from .pgvector.config import _pgvector_case_config
+            return _pgvector_case_config.get(index_type)
 
         if self == DB.PgVectoRS:
             from .pgvecto_rs.config import _pgvecto_rs_case_config

vectordb_bench/backend/clients/api.py

@@ -16,6 +16,7 @@ class IndexType(str, Enum):
     HNSW = "HNSW"
     DISKANN = "DISKANN"
     IVFFlat = "IVF_FLAT"
+    IVFSQ8 = "IVF_SQ8"
     Flat = "FLAT"
     AUTOINDEX = "AUTOINDEX"
     ES_HNSW = "hnsw"

vectordb_bench/backend/clients/chroma/chroma.py

@@ -106,21 +106,9 @@ class ChromaClient(VectorDB):
         """
         if filters:
             # assumes benchmark test filters of format: {'metadata': '>=10000', 'id': 10000}
-            metadata_value = filters.get("metadata")
             id_value = filters.get("id")
-            if metadata_value and id_value:
-                results = self.collection.query(
-                    query_embeddings=query, n_results=k, 
-                    where={"$and": [{"id": {"$eq": id_value}}, 
-                                    {"id": {"$gt": metadata_value}}
-                                    ]}
-                )
-            elif metadata_value:
-                results = self.collection.query(query_embeddings=query, n_results=k, 
-                                                where={"id": {"$gt": metadata_value}})
-            else:
-                results = self.collection.query(query_embeddings=query, n_results=k, 
-                                                where={"id": {"$eq": id_value}})
+            results = self.collection.query(query_embeddings=query, n_results=k, 
+                                                where={"id": {"$gt": id_value}})
             #return list of id's in results
             return [int(i) for i in results.get('ids')[0]]
         results = self.collection.query(query_embeddings=query, n_results=k)

vectordb_bench/backend/clients/milvus/config.py

@@ -95,6 +95,24 @@ class IVFFlatConfig(MilvusIndexConfig, DBCaseConfig):
             "metric_type": self.parse_metric(),
             "params": {"nprobe": self.nprobe},
         }
+        
+class IVFSQ8Config(MilvusIndexConfig, DBCaseConfig):
+    nlist: int
+    nprobe: int | None = None
+    index: IndexType = IndexType.IVFSQ8
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"nlist": self.nlist},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"nprobe": self.nprobe},
+        }
 
 
 class FLATConfig(MilvusIndexConfig, DBCaseConfig):
@@ -210,6 +228,7 @@ _milvus_case_config = {
     IndexType.HNSW: HNSWConfig,
     IndexType.DISKANN: DISKANNConfig,
     IndexType.IVFFlat: IVFFlatConfig,
+    IndexType.IVFSQ8: IVFSQ8Config,
     IndexType.Flat: FLATConfig,
     IndexType.GPU_IVF_FLAT: GPUIVFFlatConfig,
     IndexType.GPU_IVF_PQ: GPUIVFPQConfig,

vectordb_bench/backend/clients/pgvecto_rs/config.py

@@ -8,42 +8,30 @@ POSTGRE_URL_PLACEHOLDER = "postgresql://%s:%s@%s/%s"
 class PgVectoRSConfig(DBConfig):
     user_name: SecretStr = "postgres"
     password: SecretStr
-    url: SecretStr
+    host: str = "localhost"
+    port: int = 5432
     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()
-        host, port = url_str.split(":")
         return {
-            "host": host,
-            "port": port,
+            "host": self.host,
+            "port": self.port,
             "dbname": self.db_name,
             "user": user_str,
-            "password": pwd_str,
+            "password": pwd_str
         }
 
-
 class PgVectoRSIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None
-    quantizationType: Literal["trivial", "scalar", "product"]
-    quantizationRatio: None | Literal["x4", "x8", "x16", "x32", "x64"]
-
-    def parse_quantization(self) -> str:
-        if self.quantizationType == "trivial":
-            return "quantization = { trivial = { } }"
-        elif self.quantizationType == "scalar":
-            return "quantization = { scalar = { } }"
-        else:
-            return f'quantization = {{ product = {{ ratio = "{self.quantizationRatio}" }} }}'
 
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
-            return "l2_ops"
+            return "vector_l2_ops"
         elif self.metric_type == MetricType.IP:
-            return "dot_ops"
-        return "cosine_ops"
+            return "vector_dot_ops"
+        return "vector_cos_ops"
 
     def parse_metric_fun_op(self) -> str:
         if self.metric_type == MetricType.L2:
@@ -52,16 +40,27 @@ class PgVectoRSIndexConfig(BaseModel, DBCaseConfig):
             return "<#>"
         return "<=>"
 
+class PgVectoRSQuantConfig(PgVectoRSIndexConfig):
+    quantizationType: Literal["trivial", "scalar", "product"]
+    quantizationRatio: None | Literal["x4", "x8", "x16", "x32", "x64"]
 
-class HNSWConfig(PgVectoRSIndexConfig):
+    def parse_quantization(self) -> str:
+        if self.quantizationType == "trivial":
+            return "quantization = { trivial = { } }"
+        elif self.quantizationType == "scalar":
+            return "quantization = { scalar = { } }"
+        else:
+            return f'quantization = {{ product = {{ ratio = "{self.quantizationRatio}" }} }}'
+
+
+class HNSWConfig(PgVectoRSQuantConfig):
     M: int
     efConstruction: int
     index: IndexType = IndexType.HNSW
 
     def index_param(self) -> dict:
         options = f"""
-capacity = 1048576
-[algorithm.hnsw]
+[indexing.hnsw]
 m = {self.M}
 ef_construction = {self.efConstruction}
 {self.parse_quantization()}
@@ -72,17 +71,16 @@ ef_construction = {self.efConstruction}
         return {"metrics_op": self.parse_metric_fun_op()}
 
 
-class IVFFlatConfig(PgVectoRSIndexConfig):
+class IVFFlatConfig(PgVectoRSQuantConfig):
     nlist: int
     nprobe: int | None = None
     index: IndexType = IndexType.IVFFlat
 
     def index_param(self) -> dict:
         options = f"""
-capacity = 1048576
-[algorithm.ivf]
+[indexing.ivf]
 nlist = {self.nlist}
-nprob = {self.nprobe if self.nprobe else 10}
+nsample = {self.nprobe if self.nprobe else 10}
 {self.parse_quantization()}
 """
         return {"options": options, "metric": self.parse_metric()}
@@ -90,14 +88,29 @@ nprob = {self.nprobe if self.nprobe else 10}
     def search_param(self) -> dict:
         return {"metrics_op": self.parse_metric_fun_op()}
 
+class IVFFlatSQ8Config(PgVectoRSIndexConfig):
+    nlist: int
+    nprobe: int | None = None
+    index: IndexType = IndexType.IVFSQ8
+
+    def index_param(self) -> dict:
+        options = f"""
+[indexing.ivf]
+nlist = {self.nlist}
+nsample = {self.nprobe if self.nprobe else 10}
+quantization = {{ scalar = {{ }} }}
+"""
+        return {"options": options, "metric": self.parse_metric()}
+
+    def search_param(self) -> dict:
+        return {"metrics_op": self.parse_metric_fun_op()}
 
-class FLATConfig(PgVectoRSIndexConfig):
+class FLATConfig(PgVectoRSQuantConfig):
     index: IndexType = IndexType.Flat
 
     def index_param(self) -> dict:
         options = f"""
-capacity = 1048576
-[algorithm.flat]
+[indexing.flat]
 {self.parse_quantization()}
 """
         return {"options": options, "metric": self.parse_metric()}
@@ -107,9 +120,8 @@ capacity = 1048576
 
 
 _pgvecto_rs_case_config = {
-    IndexType.AUTOINDEX: HNSWConfig,
     IndexType.HNSW: HNSWConfig,
-    IndexType.DISKANN: HNSWConfig,
     IndexType.IVFFlat: IVFFlatConfig,
+    IndexType.IVFSQ8: IVFFlatSQ8Config,
     IndexType.Flat: FLATConfig,
 }

vectordb_bench/backend/clients/pgvecto_rs/pgvecto_rs.py

@@ -1,18 +1,17 @@
-"""Wrapper around the Pgvector vector database over VectorDB"""
+"""Wrapper around the Pgvecto.rs vector database over VectorDB"""
 
 import io
 import logging
 from contextlib import contextmanager
 from typing import Any
 import pandas as pd
-
 import psycopg2
+import psycopg2.extras
 
 from ..api import VectorDB, DBCaseConfig
 
 log = logging.getLogger(__name__)
 
-
 class PgVectoRS(VectorDB):
     """Use SQLAlchemy instructions"""
 
@@ -66,6 +65,8 @@ class PgVectoRS(VectorDB):
         self.conn = psycopg2.connect(**self.db_config)
         self.conn.autocommit = False
         self.cursor = self.conn.cursor()
+        self.cursor.execute('SET search_path = "$user", public, vectors')
+        self.conn.commit()
 
         try:
             yield
@@ -113,7 +114,7 @@ class PgVectoRS(VectorDB):
             self.conn.commit()
         except Exception as e:
             log.warning(
-                f"Failed to create pgvector table: {self.table_name} error: {e}"
+                f"Failed to create pgvecto.rs table: {self.table_name} error: {e}"
             )
             raise e from None
 
@@ -127,13 +128,10 @@ class PgVectoRS(VectorDB):
                 f'CREATE TABLE IF NOT EXISTS public."{self.table_name}" \
                     (id Integer PRIMARY KEY, embedding vector({dim}));'
             )
-            self.cursor.execute(
-                f'ALTER TABLE public."{self.table_name}" ALTER COLUMN embedding SET STORAGE PLAIN;'
-            )
             self.conn.commit()
         except Exception as e:
             log.warning(
-                f"Failed to create pgvector table: {self.table_name} error: {e}"
+                f"Failed to create pgvecto.rs table: {self.table_name} error: {e}"
             )
             raise e from None
 
@@ -146,22 +144,24 @@ class PgVectoRS(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
         try:
-            items = {"id": metadata, "embedding": embeddings}
+            items = {
+                "id": metadata,
+                "embedding": embeddings
+            }
             df = pd.DataFrame(items)
             csv_buffer = io.StringIO()
             df.to_csv(csv_buffer, index=False, header=False)
             csv_buffer.seek(0)
-            self.cursor.copy_expert(
-                f'COPY public."{self.table_name}" FROM STDIN WITH (FORMAT CSV)',
-                csv_buffer,
-            )
+            self.cursor.copy_expert(f"COPY public.\"{self.table_name}\" FROM STDIN WITH (FORMAT CSV)", csv_buffer)
             self.conn.commit()
             return len(metadata), None
         except Exception as e:
-            log.warning(
-                f"Failed to insert data into pgvector table ({self.table_name}), error: {e}"
-            )
+            log.warning(f"Failed to insert data into pgvecto.rs table ({self.table_name}), error: {e}")
+            return 0, e
 
     def search_embedding(
         self,

vectordb_bench/backend/clients/pgvector/config.py

@@ -1,49 +1,100 @@
 from pydantic import BaseModel, SecretStr
-from ..api import DBConfig, DBCaseConfig, MetricType
+from ..api import DBConfig, DBCaseConfig, IndexType, MetricType
 
 POSTGRE_URL_PLACEHOLDER = "postgresql://%s:%s@%s/%s"
 
 class PgVectorConfig(DBConfig):
     user_name: SecretStr = "postgres"
     password: SecretStr
-    url: SecretStr
+    host: str = "localhost"
+    port: int = 5432
     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)
+            "host" : self.host,
+            "port" : self.port,
+            "dbname" : self.db_name,
+            "user" : user_str,
+            "password" : pwd_str
         }
 
 class PgVectorIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None
-    lists: int | None = 1000
-    probes: int | None = 10
+    index: IndexType
 
-    def parse_metric(self) -> str: 
+    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: 
+
+    def parse_metric_fun_op(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "<->"
+        elif self.metric_type == MetricType.IP:
+            return "<#>"
+        return "<=>"
+
+    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"
 
+
+
+class HNSWConfig(PgVectorIndexConfig):
+    M: int
+    efConstruction: int
+    ef: int | None = None
+    index: IndexType = IndexType.HNSW
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"M": self.M, "efConstruction": self.efConstruction},
+        }
+
+    def index_param(self) -> dict:
+        return {
+            "m" : self.M,
+            "efConstruction" : self.efConstruction,
+            "metric" : self.parse_metric()
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "ef" : self.ef,
+            "metric_fun" : self.parse_metric_fun_str(),
+            "metric_fun_op" : self.parse_metric_fun_op(),
+        }
+
+
+class IVFFlatConfig(PgVectorIndexConfig):
+    lists: int | None = 1000
+    probes: int | None = 10
+    index: IndexType = IndexType.IVFFlat
+
     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()
-        }
+            "metric_fun" : self.parse_metric_fun_str(),
+            "metric_fun_op" : self.parse_metric_fun_op(),
+        }
+
+_pgvector_case_config = {
+    IndexType.HNSW: HNSWConfig,
+    IndexType.IVFFlat: IVFFlatConfig,
+}