vectordb-bench 0.0.7__tar.gz → 0.0.9__tar.gz

{vectordb-bench-0.0.7 → vectordb_bench-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectordb-bench
-Version: 0.0.7
+Version: 0.0.9
 Summary: VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze.
 Author-email: XuanYang-cn <xuan.yang@zilliz.com>
 Project-URL: repository, https://github.com/zilliztech/VectorDBBench
@@ -12,7 +12,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: pytz
 Requires-Dist: streamlit-autorefresh
-Requires-Dist: streamlit>=1.23.0
+Requires-Dist: streamlit!=1.34.0
 Requires-Dist: streamlit_extras
 Requires-Dist: tqdm
 Requires-Dist: s3fs
@@ -39,6 +39,7 @@ Requires-Dist: sqlalchemy; extra == "all"
 Requires-Dist: redis; extra == "all"
 Requires-Dist: chromadb; extra == "all"
 Requires-Dist: psycopg2; extra == "all"
+Requires-Dist: psycopg; extra == "all"
 Provides-Extra: qdrant
 Requires-Dist: qdrant-client; extra == "qdrant"
 Provides-Extra: pinecone
@@ -49,7 +50,7 @@ Provides-Extra: elastic
 Requires-Dist: elasticsearch; extra == "elastic"
 Provides-Extra: pgvector
 Requires-Dist: pgvector; extra == "pgvector"
-Requires-Dist: sqlalchemy; extra == "pgvector"
+Requires-Dist: psycopg; extra == "pgvector"
 Provides-Extra: pgvecto-rs
 Requires-Dist: psycopg2; extra == "pgvecto-rs"
 Provides-Extra: redis

{vectordb-bench-0.0.7 → vectordb_bench-0.0.9}/pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 dependencies = [
     "pytz",
     "streamlit-autorefresh",
-    "streamlit>=1.23.0",
+    "streamlit!=1.34.0",
     "streamlit_extras",
     "tqdm",
     "s3fs",
@@ -59,13 +59,14 @@ all = [
     "redis",
     "chromadb",
     "psycopg2",
+    "psycopg",
 ]
 
 qdrant = [ "qdrant-client" ]
 pinecone = [ "pinecone-client" ]
 weaviate = [ "weaviate-client" ]
 elastic = [ "elasticsearch" ]
-pgvector = [ "pgvector", "sqlalchemy" ]
+pgvector = [ "pgvector", "psycopg" ]
 pgvecto_rs = [ "psycopg2" ]
 redis = [ "redis" ]
 chromadb = [ "chromadb" ]

{vectordb-bench-0.0.7 → vectordb_bench-0.0.9}/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-0.0.7 → vectordb_bench-0.0.9}/vectordb_bench/backend/clients/api.py

@@ -20,6 +20,7 @@ class IndexType(str, Enum):
     Flat = "FLAT"
     AUTOINDEX = "AUTOINDEX"
     ES_HNSW = "hnsw"
+    ES_IVFFlat = "ivfflat"
     GPU_IVF_FLAT = "GPU_IVF_FLAT"
     GPU_IVF_PQ = "GPU_IVF_PQ"
     GPU_CAGRA = "GPU_CAGRA"

{vectordb-bench-0.0.7 → vectordb_bench-0.0.9}/vectordb_bench/backend/clients/milvus/milvus.py

@@ -89,6 +89,7 @@ class Milvus(VectorDB):
         connections.disconnect("default")
 
     def _optimize(self):
+        self._post_insert()
         log.info(f"{self.name} optimizing before search")
         try:
             self.col.load()
@@ -116,7 +117,7 @@ class Milvus(VectorDB):
                     time.sleep(5)
 
             wait_index()
-            
+
             # Skip compaction if use GPU indexType
             if self.case_config.index in [IndexType.GPU_CAGRA, IndexType.GPU_IVF_FLAT, IndexType.GPU_IVF_PQ]:
                 log.debug("skip compaction for gpu index type.")
@@ -179,8 +180,6 @@ class Milvus(VectorDB):
                 ]
                 res = self.col.insert(insert_data)
                 insert_count += len(res.primary_keys)
-            if kwargs.get("last_batch"):
-                self._post_insert()
         except MilvusException as e:
             log.info(f"Failed to insert data: {e}")
             return (insert_count, e)

{vectordb-bench-0.0.7 → vectordb_bench-0.0.9}/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-0.0.7 → vectordb_bench-0.0.9}/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-0.0.9/vectordb_bench/backend/clients/pgvector/config.py

@@ -0,0 +1,215 @@
+from abc import abstractmethod
+from typing import Any, Mapping, Optional, Sequence, TypedDict
+from pydantic import BaseModel, SecretStr
+from typing_extensions import LiteralString
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+POSTGRE_URL_PLACEHOLDER = "postgresql://%s:%s@%s/%s"
+
+
+class PgVectorConfigDict(TypedDict):
+    """These keys will be directly used as kwargs in psycopg connection string,
+        so the names must match exactly psycopg API"""
+
+    user: str
+    password: str
+    host: str
+    port: int
+    dbname: str
+
+
+class PgVectorConfig(DBConfig):
+    user_name: SecretStr = SecretStr("postgres")
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 5432
+    db_name: str
+
+    def to_dict(self) -> PgVectorConfigDict:
+        user_str = self.user_name.get_secret_value()
+        pwd_str = self.password.get_secret_value()
+        return {
+            "host": self.host,
+            "port": self.port,
+            "dbname": self.db_name,
+            "user": user_str,
+            "password": pwd_str,
+        }
+
+
+class PgVectorIndexParam(TypedDict):
+    metric: str
+    index_type: str
+    index_creation_with_options: Sequence[dict[str, Any]]
+    maintenance_work_mem: Optional[str]
+    max_parallel_workers: Optional[int]
+
+
+class PgVectorSearchParam(TypedDict):
+    metric_fun_op: LiteralString
+
+
+class PgVectorSessionCommands(TypedDict):
+    session_options: Sequence[dict[str, Any]]
+
+
+class PgVectorIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    create_index_before_load: bool = False
+    create_index_after_load: bool = True
+
+    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_op(self) -> LiteralString:
+        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"
+
+    @abstractmethod
+    def index_param(self) -> PgVectorIndexParam:
+        ...
+
+    @abstractmethod
+    def search_param(self) -> PgVectorSearchParam:
+        ...
+
+    @abstractmethod
+    def session_param(self) -> PgVectorSessionCommands:
+        ...
+
+    @staticmethod
+    def _optionally_build_with_options(with_options: Mapping[str, Any]) -> Sequence[dict[str, Any]]:
+        """Walk through mappings, creating a List of {key1 = value} pairs. That will be used to build a where clause"""
+        options = []
+        for option_name, value in with_options.items():
+            if value is not None:
+                options.append(
+                    {
+                        "option_name": option_name,
+                        "val": str(value),
+                    }
+                )
+        return options
+
+    @staticmethod
+    def _optionally_build_set_options(
+        set_mapping: Mapping[str, Any]
+    ) -> Sequence[dict[str, Any]]:
+        """Walk through options, creating 'SET 'key1 = "value1";' commands"""
+        session_options = []
+        for setting_name, value in set_mapping.items():
+            if value:
+                session_options.append(
+                    {"parameter": {
+                            "setting_name": setting_name,
+                            "val": str(value),
+                        },
+                    }
+                )
+        return session_options
+
+
+class PgVectorIVFFlatConfig(PgVectorIndexConfig):
+    """
+    An IVFFlat index divides vectors into lists, and then searches a subset of those lists that are
+    closest to the query vector. It has faster build times and uses less memory than HNSW,
+    but has lower query performance (in terms of speed-recall tradeoff).
+
+    Three keys to achieving good recall are:
+
+    Create the index after the table has some data
+    Choose an appropriate number of lists - a good place to start is rows / 1000 for up to 1M rows and sqrt(rows) for
+    over 1M rows.
+    When querying, specify an appropriate number of probes (higher is better for recall, lower is better for speed) -
+    a good place to start is sqrt(lists)
+    """
+
+    lists: int | None
+    probes: int | None
+    index: IndexType = IndexType.ES_IVFFlat
+    maintenance_work_mem: Optional[str] = None
+    max_parallel_workers: Optional[int] = None
+
+    def index_param(self) -> PgVectorIndexParam:
+        index_parameters = {"lists": self.lists}
+        return {
+            "metric": self.parse_metric(),
+            "index_type": self.index.value,
+            "index_creation_with_options": self._optionally_build_with_options(
+                index_parameters
+            ),
+            "maintenance_work_mem": self.maintenance_work_mem,
+            "max_parallel_workers": self.max_parallel_workers,
+        }
+
+    def search_param(self) -> PgVectorSearchParam:
+        return {
+            "metric_fun_op": self.parse_metric_fun_op(),
+        }
+
+    def session_param(self) -> PgVectorSessionCommands:
+        session_parameters = {"ivfflat.probes": self.probes}
+        return {
+            "session_options": self._optionally_build_set_options(session_parameters)
+        }
+
+
+class PgVectorHNSWConfig(PgVectorIndexConfig):
+    """
+    An HNSW index creates a multilayer graph. It has better query performance than IVFFlat (in terms of
+    speed-recall tradeoff), but has slower build times and uses more memory. Also, an index can be
+    created without any data in the table since there isn't a training step like IVFFlat.
+    """
+
+    m: int | None  # DETAIL:  Valid values are between "2" and "100".
+    ef_construction: (
+        int | None
+    )  # ef_construction must be greater than or equal to 2 * m
+    ef_search: int | None
+    index: IndexType = IndexType.ES_HNSW
+    maintenance_work_mem: Optional[str] = None
+    max_parallel_workers: Optional[int] = None
+
+    def index_param(self) -> PgVectorIndexParam:
+        index_parameters = {"m": self.m, "ef_construction": self.ef_construction}
+        return {
+            "metric": self.parse_metric(),
+            "index_type": self.index.value,
+            "index_creation_with_options": self._optionally_build_with_options(
+                index_parameters
+            ),
+            "maintenance_work_mem": self.maintenance_work_mem,
+            "max_parallel_workers": self.max_parallel_workers,
+        }
+
+    def search_param(self) -> PgVectorSearchParam:
+        return {
+            "metric_fun_op": self.parse_metric_fun_op(),
+        }
+
+    def session_param(self) -> PgVectorSessionCommands:
+        session_parameters = {"hnsw.ef_search": self.ef_search}
+        return {
+            "session_options": self._optionally_build_set_options(session_parameters)
+        }
+
+
+_pgvector_case_config = {
+        IndexType.HNSW: PgVectorHNSWConfig,
+        IndexType.ES_HNSW: PgVectorHNSWConfig,
+        IndexType.IVFFlat: PgVectorIVFFlatConfig,
+}