vectordb-bench 0.0.19__py3-none-any.whl → 0.0.21__py3-none-any.whl

vectordb_bench/backend/clients/pgvecto_rs/pgvecto_rs.py

@@ -1,14 +1,14 @@
 """Wrapper around the Pgvecto.rs vector database over VectorDB"""
 
 import logging
-import pprint
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Any, Generator, Optional, Tuple
+from typing import Any
 
 import numpy as np
 import psycopg
-from psycopg import Connection, Cursor, sql
 from pgvecto_rs.psycopg import register_vector
+from psycopg import Connection, Cursor, sql
 
 from ..api import VectorDB
 from .config import PgVectoRSConfig, PgVectoRSIndexConfig
@@ -33,7 +33,6 @@ class PgVectoRS(VectorDB):
         drop_old: bool = False,
         **kwargs,
     ):
-
         self.name = "PgVectorRS"
         self.db_config = db_config
         self.case_config = db_case_config
@@ -52,13 +51,14 @@ class PgVectoRS(VectorDB):
             (
                 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 or create_index_after_load"
-            log.error(err)
-            raise RuntimeError(
-                f"{err}\n{pprint.pformat(self.db_config)}\n{pprint.pformat(self.case_config)}"
+            msg = (
+                f"{self.name} config must create an index using create_index_before_load or create_index_after_load"
+                f"{self.name} config values: {self.db_config}\n{self.case_config}"
             )
+            log.error(msg)
+            raise RuntimeError(msg)
 
         if drop_old:
             log.info(f"Pgvecto.rs client drop table : {self.table_name}")
@@ -74,7 +74,7 @@ class PgVectoRS(VectorDB):
         self.conn = None
 
     @staticmethod
-    def _create_connection(**kwargs) -> Tuple[Connection, Cursor]:
+    def _create_connection(**kwargs) -> tuple[Connection, Cursor]:
         conn = psycopg.connect(**kwargs)
 
         # create vector extension
@@ -116,21 +116,21 @@ class PgVectoRS(VectorDB):
         self._filtered_search = sql.Composed(
             [
                 sql.SQL(
-                    "SELECT id FROM public.{table_name} WHERE id >= %s ORDER BY embedding "
+                    "SELECT id FROM public.{table_name} WHERE id >= %s ORDER BY embedding ",
                 ).format(table_name=sql.Identifier(self.table_name)),
                 sql.SQL(self.case_config.search_param()["metric_fun_op"]),
                 sql.SQL(" %s::vector LIMIT %s::int"),
-            ]
+            ],
         )
 
         self._unfiltered_search = sql.Composed(
             [
-                sql.SQL(
-                    "SELECT id FROM public.{table_name} ORDER BY embedding "
-                ).format(table_name=sql.Identifier(self.table_name)),
+                sql.SQL("SELECT id FROM public.{table_name} ORDER BY embedding ").format(
+                    table_name=sql.Identifier(self.table_name),
+                ),
                 sql.SQL(self.case_config.search_param()["metric_fun_op"]),
                 sql.SQL(" %s::vector LIMIT %s::int"),
-            ]
+            ],
         )
 
         try:
@@ -148,15 +148,12 @@ class PgVectoRS(VectorDB):
 
         self.cursor.execute(
             sql.SQL("DROP TABLE IF EXISTS public.{table_name}").format(
-                table_name=sql.Identifier(self.table_name)
-            )
+                table_name=sql.Identifier(self.table_name),
+            ),
         )
         self.conn.commit()
 
-    def ready_to_load(self):
-        pass
-
-    def optimize(self):
+    def optimize(self, data_size: int | None = None):
         self._post_insert()
 
     def _post_insert(self):
@@ -171,7 +168,7 @@ class PgVectoRS(VectorDB):
         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)
+            index_name=sql.Identifier(self._index_name),
         )
         log.debug(drop_index_sql.as_string(self.cursor))
         self.cursor.execute(drop_index_sql)
@@ -186,9 +183,9 @@ class PgVectoRS(VectorDB):
 
         index_create_sql = sql.SQL(
             """
-            CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} 
+            CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name}
             USING vectors (embedding {embedding_metric}) WITH (options = {index_options})
-            """
+            """,
         ).format(
             index_name=sql.Identifier(self._index_name),
             table_name=sql.Identifier(self.table_name),
@@ -200,10 +197,7 @@ class PgVectoRS(VectorDB):
             self.cursor.execute(index_create_sql)
             self.conn.commit()
         except Exception as e:
-            log.warning(
-                f"Failed to create pgvecto.rs index {self._index_name} \
-                    at table {self.table_name} error: {e}"
-            )
+            log.warning(f"Failed to create pgvecto.rs index {self._index_name} at table {self.table_name} error: {e}")
             raise e from None
 
     def _create_table(self, dim: int):
@@ -214,7 +208,7 @@ class PgVectoRS(VectorDB):
             """
             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,
@@ -224,9 +218,7 @@ class PgVectoRS(VectorDB):
             self.cursor.execute(table_create_sql)
             self.conn.commit()
         except Exception as e:
-            log.warning(
-                f"Failed to create pgvecto.rs table: {self.table_name} error: {e}"
-            )
+            log.warning(f"Failed to create pgvecto.rs table: {self.table_name} error: {e}")
             raise e from None
 
     def insert_embeddings(
@@ -234,7 +226,7 @@ class PgVectoRS(VectorDB):
         embeddings: list[list[float]],
         metadata: list[int],
         **kwargs: Any,
-    ) -> Tuple[int, Optional[Exception]]:
+    ) -> tuple[int, Exception | None]:
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
@@ -247,8 +239,8 @@ class PgVectoRS(VectorDB):
 
             with self.cursor.copy(
                 sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format(
-                    table_name=sql.Identifier(self.table_name)
-                )
+                    table_name=sql.Identifier(self.table_name),
+                ),
             ) as copy:
                 copy.set_types(["bigint", "vector"])
                 for i, row in enumerate(metadata_arr):
@@ -260,9 +252,7 @@ class PgVectoRS(VectorDB):
 
             return len(metadata), None
         except Exception as e:
-            log.warning(
-                f"Failed to insert data into pgvecto.rs 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(
@@ -281,12 +271,13 @@ class PgVectoRS(VectorDB):
             log.debug(self._filtered_search.as_string(self.cursor))
             gt = filters.get("id")
             result = self.cursor.execute(
-                self._filtered_search, (gt, q, k), prepare=True, binary=True
+                self._filtered_search,
+                (gt, q, k),
+                prepare=True,
+                binary=True,
             )
         else:
             log.debug(self._unfiltered_search.as_string(self.cursor))
-            result = self.cursor.execute(
-                self._unfiltered_search, (q, k), prepare=True, binary=True
-            )
+            result = self.cursor.execute(self._unfiltered_search, (q, k), prepare=True, binary=True)
 
         return [int(i[0]) for i in result.fetchall()]

vectordb_bench/backend/clients/pgvector/cli.py

@@ -1,9 +1,10 @@
-from typing import Annotated, Optional, TypedDict, Unpack
+import os
+from typing import Annotated, Unpack
 
 import click
-import os
 from pydantic import SecretStr
 
+from vectordb_bench.backend.clients import DB
 from vectordb_bench.backend.clients.api import MetricType
 
 from ....cli.cli import (
@@ -15,39 +16,48 @@ from ....cli.cli import (
     get_custom_case_config,
     run,
 )
-from vectordb_bench.backend.clients import DB
 
 
-def set_default_quantized_fetch_limit(ctx, param, value):
+# ruff: noqa
+def set_default_quantized_fetch_limit(ctx: any, param: any, value: any):
     if ctx.params.get("reranking") and value is None:
         # ef_search is the default value for quantized_fetch_limit as it's bound by ef_search.
         # 100 is default value for quantized_fetch_limit for IVFFlat.
-        default_value = ctx.params["ef_search"] if ctx.command.name == "pgvectorhnsw" else 100
-        return default_value
+        return ctx.params["ef_search"] if ctx.command.name == "pgvectorhnsw" else 100
     return value
 
+
 class PgVectorTypedDict(CommonTypedDict):
     user_name: Annotated[
-        str, click.option("--user-name", type=str, help="Db username", required=True)
+        str,
+        click.option("--user-name", type=str, help="Db username", required=True),
     ]
     password: Annotated[
         str,
-        click.option("--password",
-                     type=str,
-                     help="Postgres database password",
-                     default=lambda: os.environ.get("POSTGRES_PASSWORD", ""),
-                     show_default="$POSTGRES_PASSWORD",
-                     ),
+        click.option(
+            "--password",
+            type=str,
+            help="Postgres database password",
+            default=lambda: os.environ.get("POSTGRES_PASSWORD", ""),
+            show_default="$POSTGRES_PASSWORD",
+        ),
     ]
 
-    host: Annotated[
-        str, click.option("--host", type=str, help="Db host", required=True)
-    ]
-    db_name: Annotated[
-        str, click.option("--db-name", type=str, help="Db name", required=True)
+    host: Annotated[str, click.option("--host", type=str, help="Db host", required=True)]
+    port: Annotated[
+        int,
+        click.option(
+            "--port",
+            type=int,
+            help="Postgres database port",
+            default=5432,
+            show_default=True,
+            required=False,
+        ),
     ]
+    db_name: Annotated[str, click.option("--db-name", type=str, help="Db name", required=True)]
     maintenance_work_mem: Annotated[
-        Optional[str],
+        str | None,
         click.option(
             "--maintenance-work-mem",
             type=str,
@@ -59,7 +69,7 @@ class PgVectorTypedDict(CommonTypedDict):
         ),
     ]
     max_parallel_workers: Annotated[
-        Optional[int],
+        int | None,
         click.option(
             "--max-parallel-workers",
             type=int,
@@ -68,7 +78,7 @@ class PgVectorTypedDict(CommonTypedDict):
         ),
     ]
     quantization_type: Annotated[
-        Optional[str],
+        str | None,
         click.option(
             "--quantization-type",
             type=click.Choice(["none", "bit", "halfvec"]),
@@ -77,7 +87,7 @@ class PgVectorTypedDict(CommonTypedDict):
         ),
     ]
     reranking: Annotated[
-        Optional[bool],
+        bool | None,
         click.option(
             "--reranking/--skip-reranking",
             type=bool,
@@ -86,11 +96,11 @@ class PgVectorTypedDict(CommonTypedDict):
         ),
     ]
     reranking_metric: Annotated[
-        Optional[str],
+        str | None,
         click.option(
             "--reranking-metric",
             type=click.Choice(
-                [metric.value for metric in MetricType if metric.value not in ["HAMMING", "JACCARD"]]
+                [metric.value for metric in MetricType if metric.value not in ["HAMMING", "JACCARD"]],
             ),
             help="Distance metric for reranking",
             default="COSINE",
@@ -98,7 +108,7 @@ class PgVectorTypedDict(CommonTypedDict):
         ),
     ]
     quantized_fetch_limit: Annotated[
-        Optional[int],
+        int | None,
         click.option(
             "--quantized-fetch-limit",
             type=int,
@@ -106,13 +116,11 @@ class PgVectorTypedDict(CommonTypedDict):
                 -- bound by ef_search",
             required=False,
             callback=set_default_quantized_fetch_limit,
-        )
+        ),
     ]
 
-    
 
-class PgVectorIVFFlatTypedDict(PgVectorTypedDict, IVFFlatTypedDict):
-    ...
+class PgVectorIVFFlatTypedDict(PgVectorTypedDict, IVFFlatTypedDict): ...
 
 
 @cli.command()
@@ -130,6 +138,7 @@ def PgVectorIVFFlat(
             user_name=SecretStr(parameters["user_name"]),
             password=SecretStr(parameters["password"]),
             host=parameters["host"],
+            port=parameters["port"],
             db_name=parameters["db_name"],
         ),
         db_case_config=PgVectorIVFFlatConfig(
@@ -145,8 +154,7 @@ def PgVectorIVFFlat(
     )
 
 
-class PgVectorHNSWTypedDict(PgVectorTypedDict, HNSWFlavor1):
-    ...
+class PgVectorHNSWTypedDict(PgVectorTypedDict, HNSWFlavor1): ...
 
 
 @cli.command()
@@ -164,6 +172,7 @@ def PgVectorHNSW(
             user_name=SecretStr(parameters["user_name"]),
             password=SecretStr(parameters["password"]),
             host=parameters["host"],
+            port=parameters["port"],
             db_name=parameters["db_name"],
         ),
         db_case_config=PgVectorHNSWConfig(

vectordb_bench/backend/clients/pgvector/config.py

@@ -1,7 +1,9 @@
 from abc import abstractmethod
-from typing import Any, Mapping, Optional, Sequence, TypedDict
+from collections.abc import Mapping, Sequence
+from typing import Any, LiteralString, 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"
@@ -9,7 +11,7 @@ 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"""
+    so the names must match exactly psycopg API"""
 
     user: str
     password: str
@@ -41,8 +43,8 @@ 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]
+    maintenance_work_mem: str | None
+    max_parallel_workers: int | None
 
 
 class PgVectorSearchParam(TypedDict):
@@ -59,61 +61,60 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
     create_index_after_load: bool = True
 
     def parse_metric(self) -> str:
-        if self.quantization_type == "halfvec":
-            if self.metric_type == MetricType.L2:
-                return "halfvec_l2_ops"
-            elif self.metric_type == MetricType.IP:
-                return "halfvec_ip_ops"
-            return "halfvec_cosine_ops"
-        elif self.quantization_type == "bit":
-            if self.metric_type == MetricType.JACCARD:
-                return "bit_jaccard_ops"
-            return "bit_hamming_ops"
-        else:
-            if self.metric_type == MetricType.L2:
-                return "vector_l2_ops"
-            elif self.metric_type == MetricType.IP:
-                return "vector_ip_ops"
-            return "vector_cosine_ops"
+        d = {
+            "halfvec": {
+                MetricType.L2: "halfvec_l2_ops",
+                MetricType.IP: "halfvec_ip_ops",
+                MetricType.COSINE: "halfvec_cosine_ops",
+            },
+            "bit": {
+                MetricType.JACCARD: "bit_jaccard_ops",
+                MetricType.HAMMING: "bit_hamming_ops",
+            },
+            "_fallback": {
+                MetricType.L2: "vector_l2_ops",
+                MetricType.IP: "vector_ip_ops",
+                MetricType.COSINE: "vector_cosine_ops",
+            },
+        }
+
+        if d.get(self.quantization_type) is None:
+            return d.get("_fallback").get(self.metric_type)
+        return d.get(self.quantization_type).get(self.metric_type)
 
     def parse_metric_fun_op(self) -> LiteralString:
         if self.quantization_type == "bit":
             if self.metric_type == MetricType.JACCARD:
                 return "<%>"
             return "<~>"
-        else:
-            if self.metric_type == MetricType.L2:
-                return "<->"
-            elif self.metric_type == MetricType.IP:
-                return "<#>"
-            return "<=>"
+        if self.metric_type == MetricType.L2:
+            return "<->"
+        if 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:
+        if self.metric_type == MetricType.IP:
             return "max_inner_product"
         return "cosine_distance"
-    
+
     def parse_reranking_metric_fun_op(self) -> LiteralString:
         if self.reranking_metric == MetricType.L2:
             return "<->"
-        elif self.reranking_metric == MetricType.IP:
+        if self.reranking_metric == MetricType.IP:
             return "<#>"
         return "<=>"
 
-
     @abstractmethod
-    def index_param(self) -> PgVectorIndexParam:
-        ...
+    def index_param(self) -> PgVectorIndexParam: ...
 
     @abstractmethod
-    def search_param(self) -> PgVectorSearchParam:
-        ...
+    def search_param(self) -> PgVectorSearchParam: ...
 
     @abstractmethod
-    def session_param(self) -> PgVectorSessionCommands:
-        ...
+    def session_param(self) -> PgVectorSessionCommands: ...
 
     @staticmethod
     def _optionally_build_with_options(with_options: Mapping[str, Any]) -> Sequence[dict[str, Any]]:
@@ -125,24 +126,23 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
                     {
                         "option_name": option_name,
                         "val": str(value),
-                    }
+                    },
                 )
         return options
 
     @staticmethod
-    def _optionally_build_set_options(
-        set_mapping: Mapping[str, Any]
-    ) -> Sequence[dict[str, Any]]:
+    def _optionally_build_set_options(set_mapping: Mapping[str, Any]) -> Sequence[dict[str, Any]]:
         """Walk through options, creating 'SET 'key1 = "value1";' list"""
         session_options = []
         for setting_name, value in set_mapping.items():
             if value:
                 session_options.append(
-                    {"parameter": {
+                    {
+                        "parameter": {
                             "setting_name": setting_name,
                             "val": str(value),
                         },
-                    }
+                    },
                 )
         return session_options
 
@@ -165,12 +165,12 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig):
     lists: int | None
     probes: int | None
     index: IndexType = IndexType.ES_IVFFlat
-    maintenance_work_mem: Optional[str] = None
-    max_parallel_workers: Optional[int] = None
-    quantization_type: Optional[str] = None
-    reranking: Optional[bool] = None
-    quantized_fetch_limit: Optional[int] = None
-    reranking_metric: Optional[str] = None
+    maintenance_work_mem: str | None = None
+    max_parallel_workers: int | None = None
+    quantization_type: str | None = None
+    reranking: bool | None = None
+    quantized_fetch_limit: int | None = None
+    reranking_metric: str | None = None
 
     def index_param(self) -> PgVectorIndexParam:
         index_parameters = {"lists": self.lists}
@@ -179,9 +179,7 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig):
         return {
             "metric": self.parse_metric(),
             "index_type": self.index.value,
-            "index_creation_with_options": self._optionally_build_with_options(
-                index_parameters
-            ),
+            "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,
             "quantization_type": self.quantization_type,
@@ -197,9 +195,7 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig):
 
     def session_param(self) -> PgVectorSessionCommands:
         session_parameters = {"ivfflat.probes": self.probes}
-        return {
-            "session_options": self._optionally_build_set_options(session_parameters)
-        }
+        return {"session_options": self._optionally_build_set_options(session_parameters)}
 
 
 class PgVectorHNSWConfig(PgVectorIndexConfig):
@@ -210,17 +206,15 @@ class PgVectorHNSWConfig(PgVectorIndexConfig):
     """
 
     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_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
-    quantization_type: Optional[str] = None
-    reranking: Optional[bool] = None
-    quantized_fetch_limit: Optional[int] = None
-    reranking_metric: Optional[str] = None
+    maintenance_work_mem: str | None = None
+    max_parallel_workers: int | None = None
+    quantization_type: str | None = None
+    reranking: bool | None = None
+    quantized_fetch_limit: int | None = None
+    reranking_metric: str | None = None
 
     def index_param(self) -> PgVectorIndexParam:
         index_parameters = {"m": self.m, "ef_construction": self.ef_construction}
@@ -229,9 +223,7 @@ class PgVectorHNSWConfig(PgVectorIndexConfig):
         return {
             "metric": self.parse_metric(),
             "index_type": self.index.value,
-            "index_creation_with_options": self._optionally_build_with_options(
-                index_parameters
-            ),
+            "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,
             "quantization_type": self.quantization_type,
@@ -247,13 +239,11 @@ class PgVectorHNSWConfig(PgVectorIndexConfig):
 
     def session_param(self) -> PgVectorSessionCommands:
         session_parameters = {"hnsw.ef_search": self.ef_search}
-        return {
-            "session_options": self._optionally_build_set_options(session_parameters)
-        }
+        return {"session_options": self._optionally_build_set_options(session_parameters)}
 
 
 _pgvector_case_config = {
-        IndexType.HNSW: PgVectorHNSWConfig,
-        IndexType.ES_HNSW: PgVectorHNSWConfig,
-        IndexType.IVFFlat: PgVectorIVFFlatConfig,
+    IndexType.HNSW: PgVectorHNSWConfig,
+    IndexType.ES_HNSW: PgVectorHNSWConfig,
+    IndexType.IVFFlat: PgVectorIVFFlatConfig,
 }