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

vectordb_bench/backend/clients/pgdiskann/pgdiskann.py

@@ -1,9 +1,9 @@
 """Wrapper around the pg_diskann 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
@@ -44,20 +44,21 @@ class PgDiskANN(VectorDB):
         self._primary_field = "id"
         self._vector_field = "embedding"
 
-        self.conn, self.cursor = self._create_connection(**self.db_config)        
+        self.conn, self.cursor = self._create_connection(**self.db_config)
 
         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 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:
             self._drop_index()
@@ -72,7 +73,7 @@ class PgDiskANN(VectorDB):
         self.conn = None
 
     @staticmethod
-    def _create_connection(**kwargs) -> Tuple[Connection, Cursor]:
+    def _create_connection(**kwargs) -> tuple[Connection, Cursor]:
         conn = psycopg.connect(**kwargs)
         conn.cursor().execute("CREATE EXTENSION IF NOT EXISTS pg_diskann CASCADE")
         conn.commit()
@@ -101,25 +102,25 @@ class PgDiskANN(VectorDB):
                 log.debug(command.as_string(self.cursor))
                 self.cursor.execute(command)
             self.conn.commit()
-        
+
         self._filtered_search = sql.Composed(
             [
                 sql.SQL(
-                    "SELECT id FROM public.{table_name} WHERE id >= %s ORDER BY embedding "
-                    ).format(table_name=sql.Identifier(self.table_name)),
+                    "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.{} ORDER BY embedding ").format(
-                    sql.Identifier(self.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:
@@ -137,8 +138,8 @@ class PgDiskANN(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()
 
@@ -160,7 +161,7 @@ class PgDiskANN(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)
@@ -175,64 +176,53 @@ class PgDiskANN(VectorDB):
         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"]
-                )
+                    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"]
-                )
+                    index_param["max_parallel_workers"],
+                ),
             )
             self.cursor.execute(
-                sql.SQL(
-                    "ALTER USER {} SET max_parallel_maintenance_workers TO '{}';"
-                ).format(
+                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"]
-                )
+                    index_param["max_parallel_workers"],
+                ),
             )
             self.cursor.execute(
-                sql.SQL(
-                    "ALTER USER {} SET max_parallel_workers TO '{}';"
-                ).format(
+                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.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()
-        )
+        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"
@@ -248,28 +238,23 @@ class PgDiskANN(VectorDB):
                     sql.SQL("{option_name} = {val}").format(
                         option_name=sql.Identifier(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(())
+
+        with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options)) if any(options) else sql.Composed(())
 
         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 {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"].lower()),
             embedding_metric=sql.Identifier(index_param["metric"]),
         )
-        index_create_sql_with_with_clause = (
-            index_create_sql + with_clause
-        ).join(" ")
+        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()
@@ -283,14 +268,12 @@ class PgDiskANN(VectorDB):
 
             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)
+                    "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.conn.commit()
         except Exception as e:
-            log.warning(
-                f"Failed to create pgdiskann table: {self.table_name} error: {e}"
-            )
+            log.warning(f"Failed to create pgdiskann table: {self.table_name} error: {e}")
             raise e from None
 
     def insert_embeddings(
@@ -298,7 +281,7 @@ class PgDiskANN(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"
 
@@ -308,8 +291,8 @@ class PgDiskANN(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):
@@ -321,9 +304,7 @@ class PgDiskANN(VectorDB):
 
             return len(metadata), None
         except Exception as e:
-            log.warning(
-                f"Failed to insert data into table ({self.table_name}), error: {e}"
-            )
+            log.warning(f"Failed to insert data into table ({self.table_name}), error: {e}")
             return 0, e
 
     def search_embedding(
@@ -340,11 +321,12 @@ class PgDiskANN(VectorDB):
         if filters:
             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:
-            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/pgvecto_rs/cli.py

@@ -1,9 +1,11 @@
-from typing import Annotated, Optional, Unpack
+import os
+from typing import Annotated, Unpack
 
 import click
-import os
 from pydantic import SecretStr
 
+from vectordb_bench.backend.clients import DB
+
 from ....cli.cli import (
     CommonTypedDict,
     HNSWFlavor1,
@@ -12,12 +14,12 @@ from ....cli.cli import (
     click_parameter_decorators_from_typed_dict,
     run,
 )
-from vectordb_bench.backend.clients import DB
 
 
 class PgVectoRSTypedDict(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,
@@ -30,14 +32,10 @@ class PgVectoRSTypedDict(CommonTypedDict):
         ),
     ]
 
-    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)]
+    db_name: Annotated[str, click.option("--db-name", type=str, help="Db name", required=True)]
     max_parallel_workers: Annotated[
-        Optional[int],
+        int | None,
         click.option(
             "--max-parallel-workers",
             type=int,

vectordb_bench/backend/clients/pgvecto_rs/config.py

@@ -1,11 +1,11 @@
 from abc import abstractmethod
 from typing import TypedDict
 
+from pgvecto_rs.types import Flat, Hnsw, IndexOption, Ivf, Quantization
+from pgvecto_rs.types.index import QuantizationRatio, QuantizationType
 from pydantic import BaseModel, SecretStr
-from pgvecto_rs.types import IndexOption, Ivf, Hnsw, Flat, Quantization
-from pgvecto_rs.types.index import QuantizationType, QuantizationRatio
 
-from ..api import DBConfig, DBCaseConfig, IndexType, MetricType
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
 
 POSTGRE_URL_PLACEHOLDER = "postgresql://%s:%s@%s/%s"
 
@@ -52,14 +52,14 @@ class PgVectoRSIndexConfig(BaseModel, DBCaseConfig):
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
             return "vector_l2_ops"
-        elif self.metric_type == MetricType.IP:
+        if self.metric_type == MetricType.IP:
             return "vector_dot_ops"
         return "vector_cos_ops"
 
     def parse_metric_fun_op(self) -> str:
         if self.metric_type == MetricType.L2:
             return "<->"
-        elif self.metric_type == MetricType.IP:
+        if self.metric_type == MetricType.IP:
             return "<#>"
         return "<=>"
 
@@ -85,9 +85,7 @@ class PgVectoRSHNSWConfig(PgVectoRSIndexConfig):
         if self.quantization_type is None:
             quantization = None
         else:
-            quantization = Quantization(
-                typ=self.quantization_type, ratio=self.quantization_ratio
-            )
+            quantization = Quantization(typ=self.quantization_type, ratio=self.quantization_ratio)
 
         option = IndexOption(
             index=Hnsw(
@@ -115,9 +113,7 @@ class PgVectoRSIVFFlatConfig(PgVectoRSIndexConfig):
         if self.quantization_type is None:
             quantization = None
         else:
-            quantization = Quantization(
-                typ=self.quantization_type, ratio=self.quantization_ratio
-            )
+            quantization = Quantization(typ=self.quantization_type, ratio=self.quantization_ratio)
 
         option = IndexOption(
             index=Ivf(nlist=self.lists, quantization=quantization),
@@ -139,9 +135,7 @@ class PgVectoRSFLATConfig(PgVectoRSIndexConfig):
         if self.quantization_type is None:
             quantization = None
         else:
-            quantization = Quantization(
-                typ=self.quantization_type, ratio=self.quantization_ratio
-            )
+            quantization = Quantization(typ=self.quantization_type, ratio=self.quantization_ratio)
 
         option = IndexOption(
             index=Flat(

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,8 +148,8 @@ 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()
 
@@ -171,7 +171,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 +186,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),
@@ -202,7 +202,7 @@ class PgVectoRS(VectorDB):
         except Exception as e:
             log.warning(
                 f"Failed to create pgvecto.rs index {self._index_name} \
-                    at table {self.table_name} error: {e}"
+                    at table {self.table_name} error: {e}",
             )
             raise e from None
 
@@ -214,7 +214,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 +224,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 +232,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 +245,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):
@@ -261,7 +259,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}"
+                f"Failed to insert data into pgvecto.rs table ({self.table_name}), error: {e}",
             )
             return 0, e
 
@@ -281,12 +279,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()]