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

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -11,7 +11,7 @@ from pgvector.psycopg import register_vector
 from psycopg import Connection, Cursor, sql
 
 from ..api import VectorDB
-from .config import PgVectorConfigDict, PgVectorIndexConfig
+from .config import PgVectorConfigDict, PgVectorIndexConfig, PgVectorHNSWConfig
 
 log = logging.getLogger(__name__)
 
@@ -22,7 +22,7 @@ class PgVector(VectorDB):
     conn: psycopg.Connection[Any] | None = None
     cursor: psycopg.Cursor[Any] | None = None
 
-    # TODO add filters support
+    _filtered_search: sql.Composed
     _unfiltered_search: sql.Composed
 
     def __init__(
@@ -87,6 +87,92 @@ class PgVector(VectorDB):
         assert cursor is not None, "Cursor is not initialized"
 
         return conn, cursor
+    
+    def _generate_search_query(self, filtered: bool=False) -> sql.Composed:
+        index_param = self.case_config.index_param()
+        reranking = self.case_config.search_param()["reranking"]
+        column_name = (
+            sql.SQL("binary_quantize({0})").format(sql.Identifier("embedding"))
+            if index_param["quantization_type"] == "bit"
+            else sql.SQL("embedding")
+        )
+        search_vector = (
+            sql.SQL("binary_quantize({0})").format(sql.Placeholder())
+            if index_param["quantization_type"] == "bit"
+            else sql.Placeholder()
+        )
+
+        # The following sections assume that the quantization_type value matches the quantization function name
+        if index_param["quantization_type"] != None:
+            if index_param["quantization_type"] == "bit" and reranking:
+                # Embeddings needs to be passed to binary_quantize function if quantization_type is bit
+                search_query = sql.Composed(
+                    [
+                        sql.SQL(
+                            """
+                            SELECT i.id 
+                            FROM (
+                                SELECT id, embedding {reranking_metric_fun_op} %s::vector AS distance 
+                                FROM public.{table_name} {where_clause}
+                                ORDER BY {column_name}::{quantization_type}({dim})
+                            """
+                        ).format(
+                            table_name=sql.Identifier(self.table_name),
+                            column_name=column_name,
+                            reranking_metric_fun_op=sql.SQL(self.case_config.search_param()["reranking_metric_fun_op"]),
+                            quantization_type=sql.SQL(index_param["quantization_type"]),
+                            dim=sql.Literal(self.dim),
+                            where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
+                        ),
+                        sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                        sql.SQL(
+                            """
+                                {search_vector} 
+                                LIMIT {quantized_fetch_limit}
+                            ) i
+                            ORDER BY i.distance 
+                            LIMIT %s::int
+                            """
+                        ).format(
+                            search_vector=search_vector,
+                            quantized_fetch_limit=sql.Literal(
+                                self.case_config.search_param()["quantized_fetch_limit"]
+                            ),
+                        ),
+                    ]
+                )
+            else:
+                search_query = sql.Composed(
+                    [
+                        sql.SQL(
+                            "SELECT id FROM public.{table_name} {where_clause} ORDER BY {column_name}::{quantization_type}({dim}) "
+                        ).format(
+                            table_name=sql.Identifier(self.table_name),
+                            column_name=column_name,
+                            quantization_type=sql.SQL(index_param["quantization_type"]),
+                            dim=sql.Literal(self.dim),
+                            where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
+                        ),
+                        sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                        sql.SQL(" {search_vector} LIMIT %s::int").format(search_vector=search_vector),
+                    ]
+                )
+        else:
+            search_query = sql.Composed(
+                [
+                    sql.SQL(
+                        "SELECT id FROM public.{table_name} {where_clause} ORDER BY embedding "
+                    ).format(
+                        table_name=sql.Identifier(self.table_name),
+                        where_clause=sql.SQL("WHERE id >= %s") if filtered else sql.SQL(""),
+                    ),
+                    sql.SQL(self.case_config.search_param()["metric_fun_op"]),
+                    sql.SQL(" %s::vector LIMIT %s::int"),
+                ]
+            )
+    
+        return search_query
+    
 
     @contextmanager
     def init(self) -> Generator[None, None, None]:
@@ -112,15 +198,8 @@ class PgVector(VectorDB):
                 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"),
-            ]
-        )
+        self._filtered_search = self._generate_search_query(filtered=True)
+        self._unfiltered_search = self._generate_search_query()
 
         try:
             yield
@@ -255,17 +334,39 @@ class PgVector(VectorDB):
         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"]),
-        )
+        if index_param["quantization_type"] != None:
+            index_create_sql = sql.SQL(
+                """
+                CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name}
+                USING {index_type} (({column_name}::{quantization_type}({dim})) {embedding_metric})
+                """
+            ).format(
+                index_name=sql.Identifier(self._index_name),
+                table_name=sql.Identifier(self.table_name),
+                column_name=(
+                    sql.SQL("binary_quantize({0})").format(sql.Identifier("embedding"))
+                    if index_param["quantization_type"] == "bit"
+                    else sql.Identifier("embedding")
+                ),
+                index_type=sql.Identifier(index_param["index_type"]),
+                # This assumes that the quantization_type value matches the quantization function name
+                quantization_type=sql.SQL(index_param["quantization_type"]),
+                dim=self.dim,
+                embedding_metric=sql.Identifier(index_param["metric"]),
+            )
+        else:
+            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(" ")
@@ -341,10 +442,28 @@ class PgVector(VectorDB):
         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()
+        search_param = self.case_config.search_param()
         q = np.asarray(query)
-        # TODO add filters support
-        result = self.cursor.execute(
-            self._unfiltered_search, (q, k), prepare=True, binary=True
-        )
+        if filters:
+            gt = filters.get("id")
+            if index_param["quantization_type"] == "bit" and search_param["reranking"]:
+                result = self.cursor.execute(
+                    self._filtered_search, (q, gt, q, k), prepare=True, binary=True
+                )
+            else:
+                result = self.cursor.execute(
+                    self._filtered_search, (gt, q, k), prepare=True, binary=True
+                )
+                
+        else:
+            if index_param["quantization_type"] == "bit" and search_param["reranking"]:
+                result = self.cursor.execute(
+                    self._unfiltered_search, (q, q, k), prepare=True, binary=True
+                )
+            else:
+                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/pgvectorscale/cli.py

@@ -0,0 +1,108 @@
+import click
+import os
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from typing import Annotated, Unpack
+from vectordb_bench.backend.clients import DB
+
+
+class PgVectorScaleTypedDict(CommonTypedDict):
+    user_name: Annotated[
+        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",
+                     ),
+    ]
+
+    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)
+    ]
+
+
+class PgVectorScaleDiskAnnTypedDict(PgVectorScaleTypedDict):
+    storage_layout: Annotated[
+        str,
+        click.option(
+            "--storage-layout", type=str, help="Streaming DiskANN storage layout",
+        ),
+    ]
+    num_neighbors: Annotated[
+        int,
+        click.option(
+            "--num-neighbors", type=int, help="Streaming DiskANN num neighbors",
+        ),
+    ]
+    search_list_size: Annotated[
+        int,
+        click.option(
+            "--search-list-size", type=int, help="Streaming DiskANN search list size",
+        ),
+    ]
+    max_alpha: Annotated[
+        float,
+        click.option(
+            "--max-alpha", type=float, help="Streaming DiskANN max alpha",
+        ),
+    ]
+    num_dimensions: Annotated[
+        int,
+        click.option(
+            "--num-dimensions", type=int, help="Streaming DiskANN num dimensions",
+        ),
+    ]
+    query_search_list_size: Annotated[
+        int,
+        click.option(
+            "--query-search-list-size", type=int, help="Streaming DiskANN query search list size",
+        ),
+    ]
+    query_rescore: Annotated[
+        int,
+        click.option(
+            "--query-rescore", type=int, help="Streaming DiskANN query rescore",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectorScaleDiskAnnTypedDict)
+def PgVectorScaleDiskAnn(
+    **parameters: Unpack[PgVectorScaleDiskAnnTypedDict],
+):
+    from .config import PgVectorScaleConfig, PgVectorScaleStreamingDiskANNConfig
+
+    run(
+        db=DB.PgVectorScale,
+        db_config=PgVectorScaleConfig(
+            db_label=parameters["db_label"],
+            user_name=SecretStr(parameters["user_name"]),
+            password=SecretStr(parameters["password"]),
+            host=parameters["host"],
+            db_name=parameters["db_name"],
+        ),
+        db_case_config=PgVectorScaleStreamingDiskANNConfig(
+            storage_layout=parameters["storage_layout"],
+            num_neighbors=parameters["num_neighbors"],
+            search_list_size=parameters["search_list_size"],
+            max_alpha=parameters["max_alpha"],
+            num_dimensions=parameters["num_dimensions"],
+            query_search_list_size=parameters["query_search_list_size"],
+            query_rescore=parameters["query_rescore"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/pgvectorscale/pgvectorscale.py

@@ -22,6 +22,9 @@ class PgVectorScale(VectorDB):
     conn: psycopg.Connection[Any] | None = None
     coursor: psycopg.Cursor[Any] | None = None
 
+    _unfiltered_search: sql.Composed
+    _filtered_search: sql.Composed
+
     def __init__(
         self,
         dim: int,
@@ -99,6 +102,16 @@ class PgVectorScale(VectorDB):
                 self.cursor.execute(command)
             self.conn.commit()
         
+        self._filtered_search = sql.Composed(
+            [
+                sql.SQL("SELECT id FROM public.{} WHERE id >= %s 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")
+            ]
+        )
+        
         self._unfiltered_search = sql.Composed(
             [
                 sql.SQL("SELECT id FROM public.{} ORDER BY embedding ").format(
@@ -264,9 +277,14 @@ class PgVectorScale(VectorDB):
         assert self.cursor is not None, "Cursor is not initialized"
 
         q = np.asarray(query)
-        # TODO add filters support
-        result = self.cursor.execute(
-            self._unfiltered_search, (q, k), prepare=True, binary=True
-        )
+        if filters:
+            gt = filters.get("id")
+            result = self.cursor.execute(
+                self._filtered_search, (gt, q, k), prepare=True, binary=True
+            )
+        else:
+            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/pinecone/config.py

@@ -4,12 +4,10 @@ from ..api import DBConfig
 
 class PineconeConfig(DBConfig):
     api_key: SecretStr
-    environment: SecretStr
     index_name: str
 
     def to_dict(self) -> dict:
         return {
             "api_key": self.api_key.get_secret_value(),
-            "environment": self.environment.get_secret_value(),
             "index_name": self.index_name,
         }

vectordb_bench/backend/clients/pinecone/pinecone.py

@@ -3,7 +3,7 @@
 import logging
 from contextlib import contextmanager
 from typing import Type
-
+import pinecone
 from ..api import VectorDB, DBConfig, DBCaseConfig, EmptyDBCaseConfig, IndexType
 from .config import PineconeConfig
 
@@ -11,7 +11,8 @@ from .config import PineconeConfig
 log = logging.getLogger(__name__)
 
 PINECONE_MAX_NUM_PER_BATCH = 1000
-PINECONE_MAX_SIZE_PER_BATCH = 2 * 1024 * 1024 # 2MB
+PINECONE_MAX_SIZE_PER_BATCH = 2 * 1024 * 1024  # 2MB
+
 
 class Pinecone(VectorDB):
     def __init__(
@@ -23,30 +24,25 @@ class Pinecone(VectorDB):
         **kwargs,
     ):
         """Initialize wrapper around the milvus vector database."""
-        self.index_name = db_config["index_name"]
-        self.api_key = db_config["api_key"]
-        self.environment = db_config["environment"]
-        self.batch_size = int(min(PINECONE_MAX_SIZE_PER_BATCH / (dim * 5), PINECONE_MAX_NUM_PER_BATCH))
-        # Pincone will make connections with server while import
-        # so place the import here.
-        import pinecone
-        pinecone.init(
-            api_key=self.api_key, environment=self.environment)
+        self.index_name = db_config.get("index_name", "")
+        self.api_key = db_config.get("api_key", "")
+        self.batch_size = int(
+            min(PINECONE_MAX_SIZE_PER_BATCH / (dim * 5), PINECONE_MAX_NUM_PER_BATCH)
+        )
+
+        pc = pinecone.Pinecone(api_key=self.api_key)
+        index = pc.Index(self.index_name)
+
         if drop_old:
-            list_indexes = pinecone.list_indexes()
-            if self.index_name in list_indexes:
-                index = pinecone.Index(self.index_name)
-                index_dim = index.describe_index_stats()["dimension"]
-                if (index_dim != dim):
-                    raise ValueError(
-                        f"Pinecone index {self.index_name} dimension mismatch, expected {index_dim} got {dim}")
-                log.info(
-                    f"Pinecone client delete old index: {self.index_name}")
-                index.delete(delete_all=True)
-                index.close()
-            else:
+            index_stats = index.describe_index_stats()
+            index_dim = index_stats["dimension"]
+            if index_dim != dim:
                 raise ValueError(
-                    f"Pinecone index {self.index_name} does not exist")
+                    f"Pinecone index {self.index_name} dimension mismatch, expected {index_dim} got {dim}"
+                )
+            for namespace in index_stats["namespaces"]:
+                log.info(f"Pinecone index delete namespace: {namespace}")
+                index.delete(delete_all=True, namespace=namespace)
 
         self._metadata_key = "meta"
 
@@ -59,13 +55,10 @@ class Pinecone(VectorDB):
         return EmptyDBCaseConfig
 
     @contextmanager
-    def init(self) -> None:
-        import pinecone
-        pinecone.init(
-            api_key=self.api_key, environment=self.environment)
-        self.index = pinecone.Index(self.index_name)
+    def init(self):
+        pc = pinecone.Pinecone(api_key=self.api_key)
+        self.index = pc.Index(self.index_name)
         yield
-        self.index.close()
 
     def ready_to_load(self):
         pass
@@ -83,11 +76,16 @@ class Pinecone(VectorDB):
         insert_count = 0
         try:
             for batch_start_offset in range(0, len(embeddings), self.batch_size):
-                batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
+                batch_end_offset = min(
+                    batch_start_offset + self.batch_size, len(embeddings)
+                )
                 insert_datas = []
                 for i in range(batch_start_offset, batch_end_offset):
-                    insert_data = (str(metadata[i]), embeddings[i], {
-                                self._metadata_key: metadata[i]})
+                    insert_data = (
+                        str(metadata[i]),
+                        embeddings[i],
+                        {self._metadata_key: metadata[i]},
+                    )
                     insert_datas.append(insert_data)
                 self.index.upsert(insert_datas)
                 insert_count += batch_end_offset - batch_start_offset
@@ -101,7 +99,7 @@ class Pinecone(VectorDB):
         k: int = 100,
         filters: dict | None = None,
         timeout: int | None = None,
-    ) -> list[tuple[int, float]]:
+    ) -> list[int]:
         if filters is None:
             pinecone_filters = {}
         else:
@@ -111,9 +109,9 @@ class Pinecone(VectorDB):
                 top_k=k,
                 vector=query,
                 filter=pinecone_filters,
-            )['matches']
+            )["matches"]
         except Exception as e:
             print(f"Error querying index: {e}")
             raise e
-        id_res = [int(one_res['id']) for one_res in res]
+        id_res = [int(one_res["id"]) for one_res in res]
         return id_res

vectordb_bench/backend/clients/redis/cli.py

@@ -3,6 +3,9 @@ from typing import Annotated, TypedDict, Unpack
 import click
 from pydantic import SecretStr
 
+from .config import RedisHNSWConfig
+
+
 from ....cli.cli import (
     CommonTypedDict,
     HNSWFlavor2,
@@ -69,6 +72,11 @@ def Redis(**parameters: Unpack[RedisHNSWTypedDict]):
             ssl=parameters["ssl"],
             ssl_ca_certs=parameters["ssl_ca_certs"],
             cmd=parameters["cmd"],
+        ),        
+        db_case_config=RedisHNSWConfig(
+            M=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef=parameters["ef_runtime"],
         ),
         **parameters,
     )

vectordb_bench/backend/clients/redis/config.py

@@ -1,14 +1,45 @@
-from pydantic import SecretStr
-from ..api import DBConfig
+from pydantic import SecretStr, BaseModel
+from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
 
 class RedisConfig(DBConfig):
-    password: SecretStr
+    password: SecretStr | None = None
     host: SecretStr
-    port: int = None 
+    port: int | None = None 
 
     def to_dict(self) -> dict:
         return {
             "host": self.host.get_secret_value(),
             "port": self.port,
-            "password": self.password.get_secret_value(),
-        }
+            "password": self.password.get_secret_value() if self.password is not None else None,
+        }
+        
+
+
+class RedisIndexConfig(BaseModel):
+    """Base config for milvus"""
+
+    metric_type: MetricType | None = None
+
+    def parse_metric(self) -> str:
+        if not self.metric_type:
+            return ""
+        return self.metric_type.value
+    
+class RedisHNSWConfig(RedisIndexConfig, DBCaseConfig):
+    M: int
+    efConstruction: int
+    ef: int | None = None
+    index: IndexType = IndexType.HNSW
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"M": self.M, "efConstruction": self.efConstruction},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"ef": self.ef},
+        }

vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py

@@ -23,7 +23,7 @@ class WeaviateCloud(VectorDB):
         **kwargs,
     ):
         """Initialize wrapper around the weaviate vector database."""
-        db_config.update("auth_client_secret", weaviate.AuthApiKey(api_key=db_config.get("auth_client_secret")))
+        db_config.update({"auth_client_secret": weaviate.AuthApiKey(api_key=db_config.get("auth_client_secret"))})
         self.db_config = db_config
         self.case_config = db_case_config
         self.collection_name = collection_name

vectordb_bench/backend/runner/mp_runner.py

@@ -2,6 +2,7 @@ import time
 import traceback
 import concurrent
 import multiprocessing as mp
+import random
 import logging
 from typing import Iterable
 import numpy as np
@@ -46,7 +47,7 @@ class MultiProcessingSearchRunner:
             cond.wait()
 
         with self.db.init():
-            num, idx = len(test_data), 0
+            num, idx = len(test_data), random.randint(0, len(test_data) - 1)
 
             start_time = time.perf_counter()
             count = 0