vectordb-bench 0.0.7__py3-none-any.whl → 0.0.9__py3-none-any.whl

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -1,25 +1,36 @@
 """Wrapper around the Pgvector vector database over VectorDB"""
 
-import io
 import logging
+import pprint
 from contextlib import contextmanager
-from typing import Any
-import pandas as pd
-import psycopg2
-import psycopg2.extras
+from typing import Any, Generator, Optional, Tuple, Sequence
 
-from ..api import VectorDB, DBCaseConfig
+import numpy as np
+import psycopg
+from pgvector.psycopg import register_vector
+from psycopg import Connection, Cursor, sql
+
+from ..api import VectorDB
+from .config import PgVectorConfigDict, PgVectorIndexConfig
+
+log = logging.getLogger(__name__)
 
-log = logging.getLogger(__name__) 
 
 class PgVector(VectorDB):
-    """ Use SQLAlchemy instructions"""
+    """Use psycopg instructions"""
+
+    conn: psycopg.Connection[Any] | None = None
+    cursor: psycopg.Cursor[Any] | None = None
+
+    # TODO add filters support
+    _unfiltered_search: sql.Composed
+
     def __init__(
         self,
         dim: int,
-        db_config: dict,
-        db_case_config: DBCaseConfig,
-        collection_name: str = "PgVectorCollection",
+        db_config: PgVectorConfigDict,
+        db_case_config: PgVectorIndexConfig,
+        collection_name: str = "pg_vector_collection",
         drop_old: bool = False,
         **kwargs,
     ):
@@ -29,44 +40,89 @@ class PgVector(VectorDB):
         self.table_name = collection_name
         self.dim = dim
 
-        self._index_name = "pqvector_index"
+        self._index_name = "pgvector_index"
         self._primary_field = "id"
         self._vector_field = "embedding"
 
         # construct basic units
-        self.conn = psycopg2.connect(**self.db_config)
-        self.conn.autocommit = False
-        self.cursor = self.conn.cursor() 
-        
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+
         # create vector extension
-        self.cursor.execute('CREATE EXTENSION IF NOT EXISTS vector')
+        self.cursor.execute("CREATE EXTENSION IF NOT EXISTS vector")
         self.conn.commit()
-        
-        if drop_old :
-            log.info(f"Pgvector client drop table : {self.table_name}")
+
+        log.info(f"{self.name} config values: {self.db_config}\n{self.case_config}")
+        if not any(
+            (
+                self.case_config.create_index_before_load,
+                self.case_config.create_index_after_load,
+            )
+        ):
+            err = f"{self.name} config must create an index using create_index_before_load and/or create_index_after_load"
+            log.error(err)
+            raise RuntimeError(
+                f"{err}\n{pprint.pformat(self.db_config)}\n{pprint.pformat(self.case_config)}"
+            )
+
+        if drop_old:
             # self.pg_table.drop(pg_engine, checkfirst=True)
             self._drop_index()
             self._drop_table()
             self._create_table(dim)
-            self._create_index()
-        
+            if self.case_config.create_index_before_load:
+                self._create_index()
+
         self.cursor.close()
         self.conn.close()
         self.cursor = None
         self.conn = None
 
+    @staticmethod
+    def _create_connection(**kwargs) -> Tuple[Connection, Cursor]:
+        conn = psycopg.connect(**kwargs)
+        register_vector(conn)
+        conn.autocommit = False
+        cursor = conn.cursor()
+
+        assert conn is not None, "Connection is not initialized"
+        assert cursor is not None, "Cursor is not initialized"
+
+        return conn, cursor
+
     @contextmanager
-    def init(self) -> None:
+    def init(self) -> Generator[None, None, None]:
         """
         Examples:
             >>> with self.init():
             >>>     self.insert_embeddings()
             >>>     self.search_embedding()
         """
-        self.conn = psycopg2.connect(**self.db_config)
-        self.conn.autocommit = False
-        self.cursor = self.conn.cursor()
-        
+
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+
+        # index configuration may have commands defined that we should set during each client session
+        session_options: Sequence[dict[str, Any]] = self.case_config.session_param()["session_options"]
+
+        if len(session_options) > 0:
+            for setting in session_options:
+                command = sql.SQL("SET {setting_name} " + "= {val};").format(
+                    setting_name=sql.Identifier(setting['parameter']['setting_name']),
+                    val=sql.Identifier(str(setting['parameter']['val'])),
+                )
+                log.debug(command.as_string(self.cursor))
+                self.cursor.execute(command)
+            self.conn.commit()
+
+        self._unfiltered_search = sql.Composed(
+            [
+                sql.SQL("SELECT id FROM public.{} ORDER BY embedding ").format(
+                    sql.Identifier(self.table_name)
+                ),
+                sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                sql.SQL(" %s::vector LIMIT %s::int"),
+            ]
+        )
+
         try:
             yield
         finally:
@@ -74,54 +130,170 @@ class PgVector(VectorDB):
             self.conn.close()
             self.cursor = None
             self.conn = None
-    
+
     def _drop_table(self):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
-        
-        self.cursor.execute(f'DROP TABLE IF EXISTS public."{self.table_name}"')
+        log.info(f"{self.name} client drop table : {self.table_name}")
+
+        self.cursor.execute(
+            sql.SQL("DROP TABLE IF EXISTS public.{table_name}").format(
+                table_name=sql.Identifier(self.table_name)
+            )
+        )
         self.conn.commit()
-    
+
     def ready_to_load(self):
         pass
 
     def optimize(self):
-        pass
-    
+        self._post_insert()
+
     def _post_insert(self):
         log.info(f"{self.name} post insert before optimize")
-        self._drop_index()
-        self._create_index()
+        if self.case_config.create_index_after_load:
+            self._drop_index()
+            self._create_index()
 
-    def ready_to_search(self):
-        pass
-        
     def _drop_index(self):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
-        
-        self.cursor.execute(f'DROP INDEX IF EXISTS "{self._index_name}"')
+        log.info(f"{self.name} client drop index : {self._index_name}")
+
+        drop_index_sql = sql.SQL("DROP INDEX IF EXISTS {index_name}").format(
+            index_name=sql.Identifier(self._index_name)
+        )
+        log.debug(drop_index_sql.as_string(self.cursor))
+        self.cursor.execute(drop_index_sql)
         self.conn.commit()
-    
+
+    def _set_parallel_index_build_param(self):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        index_param = self.case_config.index_param()
+
+        if index_param["maintenance_work_mem"] is not None:
+            self.cursor.execute(
+                sql.SQL("SET maintenance_work_mem TO {};").format(
+                    index_param["maintenance_work_mem"]
+                )
+            )
+            self.cursor.execute(
+                sql.SQL("ALTER USER {} SET maintenance_work_mem TO {};").format(
+                    sql.Identifier(self.db_config["user"]),
+                    index_param["maintenance_work_mem"],
+                )
+            )
+            self.conn.commit()
+
+        if index_param["max_parallel_workers"] is not None:
+            self.cursor.execute(
+                sql.SQL("SET max_parallel_maintenance_workers TO '{}';").format(
+                    index_param["max_parallel_workers"]
+                )
+            )
+            self.cursor.execute(
+                sql.SQL(
+                    "ALTER USER {} SET max_parallel_maintenance_workers TO '{}';"
+                ).format(
+                    sql.Identifier(self.db_config["user"]),
+                    index_param["max_parallel_workers"],
+                )
+            )
+            self.cursor.execute(
+                sql.SQL("SET max_parallel_workers TO '{}';").format(
+                    index_param["max_parallel_workers"]
+                )
+            )
+            self.cursor.execute(
+                sql.SQL(
+                    "ALTER USER {} SET max_parallel_workers TO '{}';"
+                ).format(
+                    sql.Identifier(self.db_config["user"]),
+                    index_param["max_parallel_workers"],
+                )
+            )
+            self.cursor.execute(
+                sql.SQL(
+                    "ALTER TABLE {} SET (parallel_workers = {});"
+                ).format(
+                    sql.Identifier(self.table_name),
+                    index_param["max_parallel_workers"],
+                )
+            )
+            self.conn.commit()
+
+        results = self.cursor.execute(
+            sql.SQL("SHOW max_parallel_maintenance_workers;")
+        ).fetchall()
+        results.extend(
+            self.cursor.execute(sql.SQL("SHOW max_parallel_workers;")).fetchall()
+        )
+        results.extend(
+            self.cursor.execute(sql.SQL("SHOW maintenance_work_mem;")).fetchall()
+        )
+        log.info(f"{self.name} parallel index creation parameters: {results}")
+
     def _create_index(self):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
-        
+        log.info(f"{self.name} client create index : {self._index_name}")
+
         index_param = self.case_config.index_param()
-        self.cursor.execute(f'CREATE INDEX IF NOT EXISTS {self._index_name} ON public."{self.table_name}" USING ivfflat (embedding {index_param["metric"]}) WITH (lists={index_param["lists"]});')
+        self._set_parallel_index_build_param()
+        options = []
+        for option in index_param["index_creation_with_options"]:
+            if option['val'] is not None:
+                options.append(
+                    sql.SQL("{option_name} = {val}").format(
+                        option_name=sql.Identifier(option['option_name']),
+                        val=sql.Identifier(str(option['val'])),
+                    )
+                )
+        if any(options):
+            with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options))
+        else:
+            with_clause = sql.Composed(())
+
+        index_create_sql = sql.SQL(
+            "CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} USING {index_type} (embedding {embedding_metric})"
+        ).format(
+            index_name=sql.Identifier(self._index_name),
+            table_name=sql.Identifier(self.table_name),
+            index_type=sql.Identifier(index_param["index_type"]),
+            embedding_metric=sql.Identifier(index_param["metric"]),
+        )
+        index_create_sql_with_with_clause = (
+            index_create_sql + with_clause
+        ).join(" ")
+        log.debug(index_create_sql_with_with_clause.as_string(self.cursor))
+        self.cursor.execute(index_create_sql_with_with_clause)
         self.conn.commit()
-        
-    def _create_table(self, dim : int):
+
+    def _create_table(self, dim: int):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
-        
+
         try:
+            log.info(f"{self.name} client create table : {self.table_name}")
+
             # create table
-            self.cursor.execute(f'CREATE TABLE IF NOT EXISTS public."{self.table_name}" (id BIGINT PRIMARY KEY, embedding vector({dim}));')
-            self.cursor.execute(f'ALTER TABLE public."{self.table_name}" ALTER COLUMN embedding SET STORAGE PLAIN;')
+            self.cursor.execute(
+                sql.SQL(
+                    "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));"
+                ).format(table_name=sql.Identifier(self.table_name), dim=dim)
+            )
+            self.cursor.execute(
+                sql.SQL(
+                    "ALTER TABLE public.{table_name} ALTER COLUMN embedding SET STORAGE PLAIN;"
+                ).format(table_name=sql.Identifier(self.table_name))
+            )
             self.conn.commit()
         except Exception as e:
-            log.warning(f"Failed to create pgvector table: {self.table_name} error: {e}")
+            log.warning(
+                f"Failed to create pgvector table: {self.table_name} error: {e}"
+            )
             raise e from None
 
     def insert_embeddings(
@@ -129,31 +301,35 @@ class PgVector(VectorDB):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> (int, Exception):
+    ) -> Tuple[int, Optional[Exception]]:
         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
-            }
-            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)
+            metadata_arr = np.array(metadata)
+            embeddings_arr = np.array(embeddings)
+
+            with self.cursor.copy(
+                sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format(
+                    table_name=sql.Identifier(self.table_name)
+                )
+            ) as copy:
+                copy.set_types(["bigint", "vector"])
+                for i, row in enumerate(metadata_arr):
+                    copy.write_row((row, embeddings_arr[i]))
             self.conn.commit()
-            
+
             if kwargs.get("last_batch"):
                 self._post_insert()
-                
+
             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 pgvector table ({self.table_name}), error: {e}"
+            )
             return 0, e
 
-    def search_embedding(        
+    def search_embedding(
         self,
         query: list[float],
         k: int = 100,
@@ -163,11 +339,9 @@ class PgVector(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
-        search_param =self.case_config.search_param()
-        self.cursor.execute(f'SET ivfflat.probes = {search_param["probes"]}')
-        self.cursor.execute(f"SELECT id FROM public.\"{self.table_name}\" ORDER BY embedding {search_param['metric_fun_op']} '{query}' LIMIT {k};")
-        self.conn.commit()
-        result = self.cursor.fetchall()
+        # TODO add filters support
+        result = self.cursor.execute(
+            self._unfiltered_search, (query, k), prepare=True, binary=True
+        )
 
-        return [int(i[0]) for i in result]
-        
+        return [int(i[0]) for i in result.fetchall()]

vectordb_bench/backend/clients/qdrant_cloud/config.py

@@ -1,18 +1,31 @@
 from pydantic import BaseModel, SecretStr
 
 from ..api import DBConfig, DBCaseConfig, MetricType
+from pydantic import validator
 
-
+# Allowing `api_key` to be left empty, to ensure compatibility with the open-source Qdrant.
 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": True,
-        }
+        api_key = self.api_key.get_secret_value()
+        if len(api_key) > 0:
+            return {
+                "url": self.url.get_secret_value(),
+                "api_key": self.api_key.get_secret_value(),
+                "prefer_grpc": True,
+            }
+        else:
+            return {"url": self.url.get_secret_value(),}
+        
+    @validator("*")
+    def not_empty_field(cls, v, field):
+        if field.name in ["api_key", "db_label"]:
+            return v
+        if isinstance(v, (str, SecretStr)) and len(v) == 0:
+            raise ValueError("Empty string!")
+        return v
 
 class QdrantIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None

vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py

@@ -43,8 +43,7 @@ class QdrantCloud(VectorDB):
         if drop_old:
             log.info(f"QdrantCloud client drop_old collection: {self.collection_name}")
             tmp_client.delete_collection(self.collection_name)
-
-        self._create_collection(dim, tmp_client)
+            self._create_collection(dim, tmp_client)
         tmp_client = None
 
     @contextmanager
@@ -110,13 +109,18 @@ class QdrantCloud(VectorDB):
     ) -> (int, Exception):
         """Insert embeddings into Milvus. should call self.init() first"""
         assert self.qdrant_client is not None
+        QDRANT_BATCH_SIZE = 500
         try:
             # TODO: counts
-            _ = self.qdrant_client.upsert(
-                collection_name=self.collection_name,
-                wait=True,
-                points=Batch(ids=metadata, payloads=[{self._primary_field: v} for v in metadata], vectors=embeddings)
-            )
+            for offset in range(0, len(embeddings), QDRANT_BATCH_SIZE):
+                vectors = embeddings[offset: offset + QDRANT_BATCH_SIZE]
+                ids = metadata[offset: offset + QDRANT_BATCH_SIZE]
+                payloads=[{self._primary_field: v} for v in ids]
+                _ = self.qdrant_client.upsert(
+                    collection_name=self.collection_name,
+                    wait=True,
+                    points=Batch(ids=ids, payloads=payloads, vectors=vectors),
+                )
         except Exception as e:
             log.info(f"Failed to insert data, {e}")
             return 0, e

vectordb_bench/backend/runner/serial_runner.py

@@ -46,11 +46,9 @@ class SerialInsertRunner:
                 del(emb_np)
                 log.debug(f"batch dataset size: {len(all_embeddings)}, {len(all_metadata)}")
 
-                last_batch = self.dataset.data.size - count == len(all_metadata)
                 insert_count, error = self.db.insert_embeddings(
                     embeddings=all_embeddings,
                     metadata=all_metadata,
-                    last_batch=last_batch,
                 )
                 if error is not None:
                     raise error

vectordb_bench/backend/task_runner.py

@@ -140,8 +140,8 @@ class CaseRunner(BaseModel):
                 )
 
             self._init_search_runner()
-            m.recall, m.serial_latency_p99 = self._serial_search()
             m.qps = self._conc_search()
+            m.recall, m.serial_latency_p99 = self._serial_search()
         except Exception as e:
             log.warning(f"Failed to run performance case, reason = {e}")
             traceback.print_exc()

vectordb_bench/frontend/components/run_test/caseSelector.py

@@ -65,25 +65,28 @@ def caseConfigSetting(st, allCaseConfigs, case, activedDbList):
                 key = "%s-%s-%s" % (db, case, config.label.value)
                 if config.inputType == InputType.Text:
                     caseConfig[config.label] = column.text_input(
-                        config.label.value,
+                        config.displayLabel if config.displayLabel else config.label.value,
                         key=key,
+                        help=config.inputHelp,
                         value=config.inputConfig["value"],
                     )
                 elif config.inputType == InputType.Option:
                     caseConfig[config.label] = column.selectbox(
-                        config.label.value,
+                        config.displayLabel if config.displayLabel else config.label.value,
                         config.inputConfig["options"],
                         key=key,
+                        help=config.inputHelp,
                     )
                 elif config.inputType == InputType.Number:
                     caseConfig[config.label] = column.number_input(
-                        config.label.value,
+                        config.displayLabel if config.displayLabel else config.label.value,
                         # format="%d",
                         step=config.inputConfig.get("step", 1),
                         min_value=config.inputConfig["min"],
                         max_value=config.inputConfig["max"],
                         key=key,
                         value=config.inputConfig["value"],
+                        help=config.inputHelp,
                     )
                 k += 1
         if k == 0: