vectordb-bench 0.0.12__py3-none-any.whl → 0.0.13__py3-none-any.whl

vectordb_bench/backend/clients/pgvecto_rs/config.py

@@ -1,30 +1,53 @@
-from typing import Literal
+from abc import abstractmethod
+from typing import TypedDict
+
 from pydantic import BaseModel, SecretStr
-from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+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
 
 POSTGRE_URL_PLACEHOLDER = "postgresql://%s:%s@%s/%s"
 
 
+class PgVectorRSConfigDict(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 PgVectoRSConfig(DBConfig):
-    user_name: SecretStr = "postgres"
+    user_name: str = "postgres"
     password: SecretStr
     host: str = "localhost"
     port: int = 5432
     db_name: str
 
     def to_dict(self) -> dict:
-        user_str = self.user_name.get_secret_value()
+        user_str = self.user_name
         pwd_str = self.password.get_secret_value()
         return {
             "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
+    create_index_before_load: bool = False
+    create_index_after_load: bool = True
+
+    max_parallel_workers: int | None = None
+    quantization_type: QuantizationType | None = None
+    quantization_ratio: QuantizationRatio | None = None
 
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:
@@ -40,88 +63,100 @@ class PgVectoRSIndexConfig(BaseModel, DBCaseConfig):
             return "<#>"
         return "<=>"
 
-class PgVectoRSQuantConfig(PgVectoRSIndexConfig):
-    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 search_param(self) -> dict:
+        return {
+            "metric_fun_op": self.parse_metric_fun_op(),
+        }
 
-class HNSWConfig(PgVectoRSQuantConfig):
-    M: int
-    efConstruction: int
-    index: IndexType = IndexType.HNSW
+    @abstractmethod
+    def index_param(self) -> dict[str, str]: ...
 
-    def index_param(self) -> dict:
-        options = f"""
-[indexing.hnsw]
-m = {self.M}
-ef_construction = {self.efConstruction}
-{self.parse_quantization()}
-"""
-        return {"options": options, "metric": self.parse_metric()}
+    @abstractmethod
+    def session_param(self) -> dict[str, str | int]: ...
 
-    def search_param(self) -> dict:
-        return {"metrics_op": self.parse_metric_fun_op()}
 
+class PgVectoRSHNSWConfig(PgVectoRSIndexConfig):
+    index: IndexType = IndexType.HNSW
+    m: int | None = None
+    ef_search: int | None
+    ef_construction: int | None = None
 
-class IVFFlatConfig(PgVectoRSQuantConfig):
-    nlist: int
-    nprobe: int | None = None
+    def index_param(self) -> dict[str, str]:
+        if self.quantization_type is None:
+            quantization = None
+        else:
+            quantization = Quantization(
+                typ=self.quantization_type, ratio=self.quantization_ratio
+            )
+
+        option = IndexOption(
+            index=Hnsw(
+                m=self.m,
+                ef_construction=self.ef_construction,
+                quantization=quantization,
+            ),
+            threads=self.max_parallel_workers,
+        )
+        return {"options": option.dumps(), "metric": self.parse_metric()}
+
+    def session_param(self) -> dict[str, str | int]:
+        session_parameters = {}
+        if self.ef_search is not None:
+            session_parameters["vectors.hnsw_ef_search"] = str(self.ef_search)
+        return session_parameters
+
+
+class PgVectoRSIVFFlatConfig(PgVectoRSIndexConfig):
     index: IndexType = IndexType.IVFFlat
+    probes: int | None
+    lists: int | None
+
+    def index_param(self) -> dict[str, str]:
+        if self.quantization_type is None:
+            quantization = None
+        else:
+            quantization = Quantization(
+                typ=self.quantization_type, ratio=self.quantization_ratio
+            )
 
-    def index_param(self) -> dict:
-        options = f"""
-[indexing.ivf]
-nlist = {self.nlist}
-nsample = {self.nprobe if self.nprobe else 10}
-{self.parse_quantization()}
-"""
-        return {"options": options, "metric": self.parse_metric()}
+        option = IndexOption(
+            index=Ivf(nlist=self.lists, quantization=quantization),
+            threads=self.max_parallel_workers,
+        )
+        return {"options": option.dumps(), "metric": self.parse_metric()}
 
-    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 session_param(self) -> dict[str, str | int]:
+        session_parameters = {}
+        if self.probes is not None:
+            session_parameters["vectors.ivf_nprobe"] = str(self.probes)
+        return session_parameters
 
-    def search_param(self) -> dict:
-        return {"metrics_op": self.parse_metric_fun_op()}
 
-class FLATConfig(PgVectoRSQuantConfig):
+class PgVectoRSFLATConfig(PgVectoRSIndexConfig):
     index: IndexType = IndexType.Flat
 
-    def index_param(self) -> dict:
-        options = f"""
-[indexing.flat]
-{self.parse_quantization()}
-"""
-        return {"options": options, "metric": self.parse_metric()}
+    def index_param(self) -> dict[str, str]:
+        if self.quantization_type is None:
+            quantization = None
+        else:
+            quantization = Quantization(
+                typ=self.quantization_type, ratio=self.quantization_ratio
+            )
 
-    def search_param(self) -> dict:
-        return {"metrics_op": self.parse_metric_fun_op()}
+        option = IndexOption(
+            index=Flat(
+                quantization=quantization,
+            ),
+            threads=self.max_parallel_workers,
+        )
+        return {"options": option.dumps(), "metric": self.parse_metric()}
+
+    def session_param(self) -> dict[str, str | int]:
+        return {}
 
 
 _pgvecto_rs_case_config = {
-    IndexType.HNSW: HNSWConfig,
-    IndexType.IVFFlat: IVFFlatConfig,
-    IndexType.IVFSQ8: IVFFlatSQ8Config,
-    IndexType.Flat: FLATConfig,
+    IndexType.HNSW: PgVectoRSHNSWConfig,
+    IndexType.IVFFlat: PgVectoRSIVFFlatConfig,
+    IndexType.Flat: PgVectoRSFLATConfig,
 }

vectordb_bench/backend/clients/pgvecto_rs/pgvecto_rs.py

@@ -1,73 +1,138 @@
 """Wrapper around the Pgvecto.rs 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
 
-from ..api import VectorDB, DBCaseConfig
+import numpy as np
+import psycopg
+from psycopg import Connection, Cursor, sql
+from pgvecto_rs.psycopg import register_vector
+
+from ..api import VectorDB
+from .config import PgVectoRSConfig, PgVectoRSIndexConfig
 
 log = logging.getLogger(__name__)
 
+
 class PgVectoRS(VectorDB):
-    """Use SQLAlchemy instructions"""
+    """Use psycopg instructions"""
+
+    conn: psycopg.Connection[Any] | None = None
+    cursor: psycopg.Cursor[Any] | None = None
+    _unfiltered_search: sql.Composed
+    _filtered_search: sql.Composed
 
     def __init__(
         self,
         dim: int,
-        db_config: dict,
-        db_case_config: DBCaseConfig,
-        collection_name: str = "PgVectorCollection",
+        db_config: PgVectoRSConfig,
+        db_case_config: PgVectoRSIndexConfig,
+        collection_name: str = "PgVectoRSCollection",
         drop_old: bool = False,
         **kwargs,
     ):
+
+        self.name = "PgVectorRS"
         self.db_config = db_config
         self.case_config = db_case_config
         self.table_name = collection_name
         self.dim = dim
 
-        self._index_name = "pqvector_index"
+        self._index_name = "pgvectors_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 vectors")
-        self.conn.commit()
+        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)}"
+            )
 
         if drop_old:
             log.info(f"Pgvecto.rs client drop table : {self.table_name}")
             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)
+
+        # create vector extension
+        conn.execute("CREATE EXTENSION IF NOT EXISTS vectors")
+        conn.commit()
+        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.cursor.execute('SET search_path = "$user", public, vectors')
+
+        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 = self.case_config.session_param()
+
+        for key, val in session_options.items():
+            command = sql.SQL("SET {setting_name} " + "= {val};").format(
+                setting_name=sql.Identifier(key),
+                val=val,
+            )
+            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)),
+                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(self.case_config.search_param()["metric_fun_op"]),
+                sql.SQL(" %s::vector LIMIT %s::int"),
+            ]
+        )
+
         try:
             yield
         finally:
@@ -79,42 +144,65 @@ class PgVectoRS(VectorDB):
     def _drop_table(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 drop table : {self.table_name}")
 
-        self.cursor.execute(f'DROP TABLE IF EXISTS public."{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 ready_to_search(self):
-        pass
+    def _post_insert(self):
+        log.info(f"{self.name} post insert before optimize")
+        if self.case_config.create_index_after_load:
+            self._drop_index()
+            self._create_index()
 
     def _drop_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 drop index : {self._index_name}")
 
-        self.cursor.execute(f'DROP INDEX IF EXISTS "{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 _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()
 
+        index_create_sql = sql.SQL(
+            """
+            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),
+            embedding_metric=sql.Identifier(index_param["metric"]),
+            index_options=index_param["options"],
+        )
         try:
-            # create table
-            self.cursor.execute(
-                f'CREATE INDEX IF NOT EXISTS {self._index_name} ON public."{self.table_name}" \
-                    USING vectors (embedding {index_param["metric"]}) WITH (options = $${index_param["options"]}$$);'
-            )
+            log.debug(index_create_sql.as_string(self.cursor))
+            self.cursor.execute(index_create_sql)
             self.conn.commit()
         except Exception as e:
             log.warning(
-                f"Failed to create pgvecto.rs table: {self.table_name} error: {e}"
+                f"Failed to create pgvecto.rs index {self._index_name} \
+                    at table {self.table_name} error: {e}"
             )
             raise e from None
 
@@ -122,12 +210,18 @@ class PgVectoRS(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         assert self.cursor is not None, "Cursor is not initialized"
 
+        table_create_sql = 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,
+        )
         try:
             # create table
-            self.cursor.execute(
-                f'CREATE TABLE IF NOT EXISTS public."{self.table_name}" \
-                    (id Integer PRIMARY KEY, embedding vector({dim}));'
-            )
+            self.cursor.execute(table_create_sql)
             self.conn.commit()
         except Exception as e:
             log.warning(
@@ -140,7 +234,7 @@ class PgVectoRS(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"
 
@@ -148,19 +242,27 @@ class PgVectoRS(VectorDB):
         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 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(
@@ -173,20 +275,18 @@ class PgVectoRS(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()
+        q = np.asarray(query)
 
         if filters:
+            log.debug(self._filtered_search.as_string(self.cursor))
             gt = filters.get("id")
-            self.cursor.execute(
-                f"SELECT id FROM (SELECT * FROM public.\"{self.table_name}\" ORDER BY embedding \
-                    {search_param['metrics_op']} '{query}' LIMIT {k}) AS X WHERE id > {gt} ;"
+            result = self.cursor.execute(
+                self._filtered_search, (gt, q, k), prepare=True, binary=True
             )
         else:
-            self.cursor.execute(
-                f"SELECT id FROM public.\"{self.table_name}\" ORDER BY embedding \
-                    {search_param['metrics_op']} '{query}' LIMIT {k};"
+            log.debug(self._unfiltered_search.as_string(self.cursor))
+            result = self.cursor.execute(
+                self._unfiltered_search, (q, k), prepare=True, binary=True
             )
-        self.conn.commit()
-        result = self.cursor.fetchall()
 
-        return [int(i[0]) for i in result]
+        return [int(i[0]) for i in result.fetchall()]

vectordb_bench/backend/clients/pgvectorscale/config.py

@@ -0,0 +1,111 @@
+from abc import abstractmethod
+from typing import 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 PgVectorScaleConfigDict(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 PgVectorScaleConfig(DBConfig):
+    user_name: SecretStr = SecretStr("postgres")
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 5432
+    db_name: str
+
+    def to_dict(self) -> PgVectorScaleConfigDict:
+        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 PgVectorScaleIndexConfig(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.COSINE:
+            return "vector_cosine_ops"
+        return ""
+        
+    def parse_metric_fun_op(self) -> LiteralString:
+        if self.metric_type == MetricType.COSINE:
+            return "<=>"
+        return ""
+
+    def parse_metric_fun_str(self) -> str:
+        if self.metric_type == MetricType.COSINE:
+            return "cosine_distance"
+        return ""
+    
+    @abstractmethod
+    def index_param(self) -> dict:
+        ...
+
+    @abstractmethod
+    def search_param(self) -> dict:
+        ...
+
+    @abstractmethod
+    def session_param(self) -> dict:
+        ...
+    
+
+class PgVectorScaleStreamingDiskANNConfig(PgVectorScaleIndexConfig):
+    index: IndexType = IndexType.STREAMING_DISKANN
+    storage_layout: str | None
+    num_neighbors: int | None
+    search_list_size: int | None
+    max_alpha: float | None
+    num_dimensions: int | None
+    num_bits_per_dimension: int | None
+    query_search_list_size: int | None
+    query_rescore: int | None
+
+    def index_param(self) -> dict:
+        return {
+            "metric": self.parse_metric(),
+            "index_type": self.index.value,
+            "options": {
+                "storage_layout": self.storage_layout,
+                "num_neighbors": self.num_neighbors,
+                "search_list_size": self.search_list_size,
+                "max_alpha": self.max_alpha,
+                "num_dimensions": self.num_dimensions,
+            },
+        }
+    
+    def search_param(self) -> dict:
+        return {
+            "metric": self.parse_metric(),
+            "metric_fun_op": self.parse_metric_fun_op(),
+        }
+    
+    def session_param(self) -> dict:
+        return {
+            "diskann.query_search_list_size": self.query_search_list_size,
+            "diskann.query_rescore": self.query_rescore,
+        }
+    
+_pgvectorscale_case_config = {
+    IndexType.STREAMING_DISKANN: PgVectorScaleStreamingDiskANNConfig,
+}