vectordb-bench 0.0.17__py3-none-any.whl → 0.0.18__py3-none-any.whl

vectordb_bench/backend/clients/__init__.py

@@ -32,6 +32,7 @@ class DB(Enum):
     PgVectoRS = "PgVectoRS"
     PgVectorScale = "PgVectorScale"
     PgDiskANN = "PgDiskANN"
+    AlloyDB = "AlloyDB"
     Redis = "Redis"
     MemoryDB = "MemoryDB"
     Chroma = "Chroma"
@@ -97,6 +98,10 @@ class DB(Enum):
         if self == DB.AWSOpenSearch:
             from .aws_opensearch.aws_opensearch import AWSOpenSearch
             return AWSOpenSearch
+        
+        if self == DB.AlloyDB:
+            from .alloydb.alloydb import AlloyDB
+            return AlloyDB
 
     @property
     def config_cls(self) -> Type[DBConfig]:
@@ -156,6 +161,10 @@ class DB(Enum):
         if self == DB.AWSOpenSearch:
             from .aws_opensearch.config import AWSOpenSearchConfig
             return AWSOpenSearchConfig
+        
+        if self == DB.AlloyDB:
+            from .alloydb.config import AlloyDBConfig
+            return AlloyDBConfig
 
     def case_config_cls(self, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
         if self == DB.Milvus:
@@ -197,6 +206,10 @@ class DB(Enum):
         if self == DB.PgDiskANN:
             from .pgdiskann.config import _pgdiskann_case_config
             return _pgdiskann_case_config.get(index_type)
+        
+        if self == DB.AlloyDB:
+            from .alloydb.config import _alloydb_case_config
+            return _alloydb_case_config.get(index_type)
 
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig

vectordb_bench/backend/clients/alloydb/alloydb.py

@@ -0,0 +1,372 @@
+"""Wrapper around the alloydb vector database over VectorDB"""
+
+import logging
+import pprint
+from contextlib import contextmanager
+from typing import Any, Generator, Optional, Tuple, Sequence
+
+import numpy as np
+import psycopg
+from pgvector.psycopg import register_vector
+from psycopg import Connection, Cursor, sql
+
+from ..api import VectorDB
+from .config import AlloyDBConfigDict, AlloyDBIndexConfig, AlloyDBScaNNConfig
+
+log = logging.getLogger(__name__)
+
+
+class AlloyDB(VectorDB):
+    """Use psycopg instructions"""
+
+    conn: psycopg.Connection[Any] | None = None
+    cursor: psycopg.Cursor[Any] | None = None
+
+    _filtered_search: sql.Composed
+    _unfiltered_search: sql.Composed
+
+    def __init__(
+        self,
+        dim: int,
+        db_config: AlloyDBConfigDict,
+        db_case_config: AlloyDBIndexConfig,
+        collection_name: str = "alloydb_collection",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.name = "AlloyDB"
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+
+        self._index_name = "alloydb_index"
+        self._primary_field = "id"
+        self._vector_field = "embedding"
+
+        # construct basic units
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+
+        # create vector extension
+        self.cursor.execute("CREATE EXTENSION IF NOT EXISTS alloydb_scann CASCADE")
+        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:
+            self._drop_index()
+            self._drop_table()
+            self._create_table(dim)
+            if self.case_config.create_index_before_load:
+                self._create_index()
+
+        self.cursor.close()
+        self.conn.close()
+        self.cursor = None
+        self.conn = None
+
+    @staticmethod
+    def _create_connection(**kwargs) -> Tuple[Connection, Cursor]:
+        conn = psycopg.connect(**kwargs)
+        register_vector(conn)
+        conn.autocommit = False
+        cursor = conn.cursor()
+
+        assert conn is not None, "Connection is not initialized"
+        assert cursor is not None, "Cursor is not initialized"
+        return conn, cursor
+    
+    def _generate_search_query(self, filtered: bool=False) -> sql.Composed:
+        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]:
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+
+        self.conn, self.cursor = self._create_connection(**self.db_config)
+
+        # index configuration may have commands defined that we should set during each client session
+        session_options: Sequence[dict[str, Any]] = self.case_config.session_param()["session_options"]
+
+        if len(session_options) > 0:
+            for setting in session_options:
+                command = sql.SQL("SET {setting_name} " + "= {val};").format(
+                    setting_name=sql.Identifier(setting['parameter']['setting_name']),
+                    val=sql.Identifier(str(setting['parameter']['val'])),
+                )
+                log.debug(command.as_string(self.cursor))
+                self.cursor.execute(command)
+            self.conn.commit()
+
+        self._filtered_search = self._generate_search_query(filtered=True)
+        self._unfiltered_search = self._generate_search_query()
+
+        try:
+            yield
+        finally:
+            self.cursor.close()
+            self.conn.close()
+            self.cursor = None
+            self.conn = None
+
+    def _drop_table(self):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+        log.info(f"{self.name} client drop table : {self.table_name}")
+
+        self.cursor.execute(
+            sql.SQL("DROP TABLE IF EXISTS public.{table_name}").format(
+                table_name=sql.Identifier(self.table_name)
+            )
+        )
+        self.conn.commit()
+
+    def ready_to_load(self):
+        pass
+
+    def optimize(self):
+        self._post_insert()
+
+    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}")
+
+        drop_index_sql = sql.SQL("DROP INDEX IF EXISTS {index_name}").format(
+            index_name=sql.Identifier(self._index_name)
+        )
+        log.debug(drop_index_sql.as_string(self.cursor))
+        self.cursor.execute(drop_index_sql)
+        self.conn.commit()
+
+    def _set_parallel_index_build_param(self):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        index_param = self.case_config.index_param()
+
+        if index_param["enable_pca"] is not None:
+            self.cursor.execute(
+                sql.SQL("SET scann.enable_pca TO {};").format(
+                    index_param["enable_pca"]
+                )
+            )
+            self.cursor.execute(
+                sql.SQL("ALTER USER {} SET scann.enable_pca TO {};").format(
+                    sql.Identifier(self.db_config["user"]),
+                    index_param["enable_pca"],
+                )
+            )
+            self.conn.commit()
+
+        if index_param["maintenance_work_mem"] is not None:
+            self.cursor.execute(
+                sql.SQL("SET maintenance_work_mem TO {};").format(
+                    index_param["maintenance_work_mem"]
+                )
+            )
+            self.cursor.execute(
+                sql.SQL("ALTER USER {} SET maintenance_work_mem TO {};").format(
+                    sql.Identifier(self.db_config["user"]),
+                    index_param["maintenance_work_mem"],
+                )
+            )
+            self.conn.commit()
+
+        if index_param["max_parallel_workers"] is not None:
+            self.cursor.execute(
+                sql.SQL("SET max_parallel_maintenance_workers TO '{}';").format(
+                    index_param["max_parallel_workers"]
+                )
+            )
+            self.cursor.execute(
+                sql.SQL(
+                    "ALTER USER {} SET max_parallel_maintenance_workers TO '{}';"
+                ).format(
+                    sql.Identifier(self.db_config["user"]),
+                    index_param["max_parallel_workers"],
+                )
+            )
+            self.cursor.execute(
+                sql.SQL("SET max_parallel_workers TO '{}';").format(
+                    index_param["max_parallel_workers"]
+                )
+            )
+            self.cursor.execute(
+                sql.SQL(
+                    "ALTER USER {} SET max_parallel_workers TO '{}';"
+                ).format(
+                    sql.Identifier(self.db_config["user"]),
+                    index_param["max_parallel_workers"],
+                )
+            )
+            self.cursor.execute(
+                sql.SQL(
+                    "ALTER TABLE {} SET (parallel_workers = {});"
+                ).format(
+                    sql.Identifier(self.table_name),
+                    index_param["max_parallel_workers"],
+                )
+            )
+            self.conn.commit()
+
+        results = self.cursor.execute(
+            sql.SQL("SHOW max_parallel_maintenance_workers;")
+        ).fetchall()
+        results.extend(
+            self.cursor.execute(sql.SQL("SHOW max_parallel_workers;")).fetchall()
+        )
+        results.extend(
+            self.cursor.execute(sql.SQL("SHOW maintenance_work_mem;")).fetchall()
+        )
+        log.info(f"{self.name} parallel index creation parameters: {results}")
+
+    def _create_index(self):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+        log.info(f"{self.name} client create index : {self._index_name}")
+
+        index_param = self.case_config.index_param()
+        self._set_parallel_index_build_param()
+        options = []
+        for option in index_param["index_creation_with_options"]:
+            if option['val'] is not None:
+                options.append(
+                    sql.SQL("{option_name} = {val}").format(
+                        option_name=sql.Identifier(option['option_name']),
+                        val=sql.Identifier(str(option['val'])),
+                    )
+                )
+        if any(options):
+            with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options))
+        else:
+            with_clause = sql.Composed(())
+
+        index_create_sql = sql.SQL(
+            """
+            CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} 
+            USING {index_type} (embedding {embedding_metric})
+            """
+        ).format(
+            index_name=sql.Identifier(self._index_name),
+            table_name=sql.Identifier(self.table_name),
+            index_type=sql.Identifier(index_param["index_type"]),
+            embedding_metric=sql.Identifier(index_param["metric"]),
+        )
+
+        index_create_sql_with_with_clause = (
+            index_create_sql + with_clause
+        ).join(" ")
+        log.debug(index_create_sql_with_with_clause.as_string(self.cursor))
+        self.cursor.execute(index_create_sql_with_with_clause)
+        self.conn.commit()
+
+    def _create_table(self, dim: int):
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        try:
+            log.info(f"{self.name} client create table : {self.table_name}")
+
+            # create table
+            self.cursor.execute(
+                sql.SQL(
+                    "CREATE TABLE IF NOT EXISTS public.{table_name} (id BIGINT PRIMARY KEY, embedding vector({dim}));"
+                ).format(table_name=sql.Identifier(self.table_name), dim=dim)
+            )
+            self.conn.commit()
+        except Exception as e:
+            log.warning(
+                f"Failed to create alloydb table: {self.table_name} error: {e}"
+            )
+            raise e from None
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> Tuple[int, Optional[Exception]]:
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        try:
+            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 alloydb table ({self.table_name}), error: {e}"
+            )
+            return 0, e
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+    ) -> list[int]:
+        assert self.conn is not None, "Connection is not initialized"
+        assert self.cursor is not None, "Cursor is not initialized"
+
+        q = np.asarray(query)
+        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/alloydb/cli.py

@@ -0,0 +1,147 @@
+from typing import Annotated, Optional, TypedDict, Unpack
+
+import click
+import os
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients.api import MetricType
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    get_custom_case_config,
+    run,
+)
+from vectordb_bench.backend.clients import DB
+
+
+class AlloyDBTypedDict(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)
+    ]
+    maintenance_work_mem: Annotated[
+        Optional[str],
+        click.option(
+            "--maintenance-work-mem",
+            type=str,
+            help="Sets the maximum memory to be used for maintenance operations (index creation). "
+            "Can be entered as string with unit like '64GB' or as an integer number of KB."
+            "This will set the parameters: max_parallel_maintenance_workers,"
+            " max_parallel_workers & table(parallel_workers)",
+            required=False,
+        ),
+    ]
+    max_parallel_workers: Annotated[
+        Optional[int],
+        click.option(
+            "--max-parallel-workers",
+            type=int,
+            help="Sets the maximum number of parallel processes per maintenance operation (index creation)",
+            required=False,
+        ),
+    ]
+
+    
+
+class AlloyDBScaNNTypedDict(AlloyDBTypedDict):
+    num_leaves: Annotated[
+        int,
+        click.option("--num-leaves", type=int, help="Number of leaves", required=True)
+    ]
+    num_leaves_to_search: Annotated[
+        int,
+        click.option("--num-leaves-to-search", type=int, help="Number of leaves to search", required=True)
+    ]
+    pre_reordering_num_neighbors: Annotated[
+        int,
+        click.option("--pre-reordering-num-neighbors", type=int, help="Pre-reordering number of neighbors", default=200)
+    ]
+    max_top_neighbors_buffer_size: Annotated[
+        int,
+        click.option("--max-top-neighbors-buffer-size", type=int, help="Maximum top neighbors buffer size", default=20_000)
+    ]
+    num_search_threads: Annotated[
+        int,
+        click.option("--num-search-threads", type=int, help="Number of search threads", default=2)
+    ]
+    max_num_prefetch_datasets: Annotated[
+        int,
+        click.option("--max-num-prefetch-datasets", type=int, help="Maximum number of prefetch datasets", default=100)
+    ]
+    quantizer: Annotated[
+        str,
+        click.option(
+            "--quantizer",
+            type=click.Choice(["SQ8", "FLAT"]),
+            help="Quantizer type",
+            default="SQ8"
+        )
+    ]
+    enable_pca: Annotated[
+        bool, click.option(
+            "--enable-pca",
+            type=click.Choice(["on", "off"]),
+            help="Enable PCA",
+            default="on"
+        )
+    ]
+    max_num_levels: Annotated[
+        int,
+        click.option(
+            "--max-num-levels",
+            type=click.Choice([1, 2]),
+            help="Maximum number of levels",
+            default=1
+        )
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(AlloyDBScaNNTypedDict)
+def AlloyDBScaNN(
+    **parameters: Unpack[AlloyDBScaNNTypedDict],
+):
+    from .config import AlloyDBConfig, AlloyDBScaNNConfig
+
+    parameters["custom_case"] = get_custom_case_config(parameters)
+    run(
+        db=DB.AlloyDB,
+        db_config=AlloyDBConfig(
+            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=AlloyDBScaNNConfig(
+            num_leaves=parameters["num_leaves"],
+            quantizer=parameters["quantizer"],
+            enable_pca=parameters["enable_pca"],
+            max_num_levels=parameters["max_num_levels"],
+            num_leaves_to_search=parameters["num_leaves_to_search"],
+            max_top_neighbors_buffer_size=parameters["max_top_neighbors_buffer_size"],
+            pre_reordering_num_neighbors=parameters["pre_reordering_num_neighbors"],
+            num_search_threads=parameters["num_search_threads"],
+            max_num_prefetch_datasets=parameters["max_num_prefetch_datasets"],
+            max_parallel_workers=parameters["max_parallel_workers"],
+            maintenance_work_mem=parameters["maintenance_work_mem"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/alloydb/config.py

@@ -0,0 +1,168 @@
+from abc import abstractmethod
+from typing import Any, Mapping, Optional, Sequence, 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 AlloyDBConfigDict(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 AlloyDBConfig(DBConfig):
+    user_name: SecretStr = SecretStr("postgres")
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 5432
+    db_name: str
+
+    def to_dict(self) -> AlloyDBConfigDict:
+        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 AlloyDBIndexParam(TypedDict):
+    metric: str
+    index_type: str
+    index_creation_with_options: Sequence[dict[str, Any]]
+    maintenance_work_mem: Optional[str]
+    max_parallel_workers: Optional[int]
+
+
+class AlloyDBSearchParam(TypedDict):
+    metric_fun_op: LiteralString
+
+
+class AlloyDBSessionCommands(TypedDict):
+    session_options: Sequence[dict[str, Any]]
+
+
+class AlloyDBIndexConfig(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.L2:
+            return "l2"
+        elif self.metric_type == MetricType.DP:
+            return "dot_product"
+        return "cosine"
+
+    def parse_metric_fun_op(self) -> LiteralString:
+        if self.metric_type == MetricType.L2:
+            return "<->"
+        elif self.metric_type == MetricType.IP:
+            return "<#>"
+        return "<=>"
+
+    @abstractmethod
+    def index_param(self) -> AlloyDBIndexParam:
+        ...
+
+    @abstractmethod
+    def search_param(self) -> AlloyDBSearchParam:
+        ...
+
+    @abstractmethod
+    def session_param(self) -> AlloyDBSessionCommands:
+        ...
+
+    @staticmethod
+    def _optionally_build_with_options(with_options: Mapping[str, Any]) -> Sequence[dict[str, Any]]:
+        """Walk through mappings, creating a List of {key1 = value} pairs. That will be used to build a where clause"""
+        options = []
+        for option_name, value in with_options.items():
+            if value is not None:
+                options.append(
+                    {
+                        "option_name": option_name,
+                        "val": str(value),
+                    }
+                )
+        return options
+
+    @staticmethod
+    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": {
+                            "setting_name": setting_name,
+                            "val": str(value),
+                        },
+                    }
+                )
+        return session_options
+
+
+class AlloyDBScaNNConfig(AlloyDBIndexConfig):
+    index: IndexType = IndexType.SCANN
+    num_leaves: int | None
+    quantizer: str | None
+    enable_pca: str | None
+    max_num_levels: int | None
+    num_leaves_to_search: int | None
+    max_top_neighbors_buffer_size: int | None
+    pre_reordering_num_neighbors: int |  None
+    num_search_threads: int  | None
+    max_num_prefetch_datasets: int | None
+    maintenance_work_mem: Optional[str] = None
+    max_parallel_workers: Optional[int] = None
+
+    def index_param(self) -> AlloyDBIndexParam:
+        index_parameters = {
+            "num_leaves": self.num_leaves, "max_num_levels": self.max_num_levels, "quantizer": self.quantizer,
+        }
+        return {
+            "metric": self.parse_metric(),
+            "index_type": self.index.value,
+            "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,
+            "enable_pca": self.enable_pca,
+        }
+
+    def search_param(self) -> AlloyDBSearchParam:
+        return {
+            "metric_fun_op": self.parse_metric_fun_op(),
+        }
+
+    def session_param(self) -> AlloyDBSessionCommands:
+        session_parameters = {
+            "scann.num_leaves_to_search": self.num_leaves_to_search,
+            "scann.max_top_neighbors_buffer_size": self.max_top_neighbors_buffer_size,
+            "scann.pre_reordering_num_neighbors": self.pre_reordering_num_neighbors,
+            "scann.num_search_threads": self.num_search_threads,
+            "scann.max_num_prefetch_datasets": self.max_num_prefetch_datasets,
+        }
+        return {
+            "session_options": self._optionally_build_set_options(session_parameters)
+        }
+
+
+_alloydb_case_config = {
+        IndexType.SCANN: AlloyDBScaNNConfig,
+}

vectordb_bench/backend/clients/api.py

@@ -10,6 +10,7 @@ class MetricType(str, Enum):
     L2 = "L2"
     COSINE = "COSINE"
     IP = "IP"
+    DP = "DP"
     HAMMING = "HAMMING"
     JACCARD = "JACCARD"
 
@@ -27,6 +28,7 @@ class IndexType(str, Enum):
     GPU_IVF_FLAT = "GPU_IVF_FLAT"
     GPU_IVF_PQ = "GPU_IVF_PQ"
     GPU_CAGRA = "GPU_CAGRA"
+    SCANN = "scann"
 
 
 class DBConfig(ABC, BaseModel):

vectordb_bench/backend/clients/milvus/milvus.py

@@ -93,7 +93,7 @@ class Milvus(VectorDB):
         self._post_insert()
         log.info(f"{self.name} optimizing before search")
         try:
-            self.col.load(refresh=True)
+            self.col.load()
         except Exception as e:
             log.warning(f"{self.name} optimize error: {e}")
             raise e from None

vectordb_bench/cli/vectordbbench.py

@@ -9,6 +9,7 @@ from ..backend.clients.weaviate_cloud.cli import Weaviate
 from ..backend.clients.zilliz_cloud.cli import ZillizAutoIndex
 from ..backend.clients.milvus.cli import MilvusAutoIndex
 from ..backend.clients.aws_opensearch.cli import AWSOpenSearch
+from ..backend.clients.alloydb.cli import AlloyDBScaNN
 
 from .cli import cli
 
@@ -24,6 +25,7 @@ cli.add_command(MilvusAutoIndex)
 cli.add_command(AWSOpenSearch)
 cli.add_command(PgVectorScaleDiskAnn)
 cli.add_command(PgDiskAnn)
+cli.add_command(AlloyDBScaNN)
 
 
 if __name__ == "__main__":

vectordb_bench/frontend/config/dbCaseConfigs.py

@@ -906,6 +906,141 @@ CaseConfigParamInput_reranking_metric_PgVector = CaseConfigInput(
     == "bit" and config.get(CaseConfigParamType.reranking, False)
 )
 
+
+CaseConfigParamInput_IndexType_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.IndexType,
+    inputHelp="Select Index Type",
+    inputType=InputType.Option,
+    inputConfig={
+        "options": [
+            IndexType.SCANN.value,
+        ],
+    },
+)
+
+CaseConfigParamInput_num_leaves_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.numLeaves,
+    displayLabel="Num Leaves",
+    inputHelp="The number of partition to apply to this index",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 1,
+        "max": 1048576,
+        "value": 200,
+    },
+)
+
+CaseConfigParamInput_quantizer_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.quantizer,
+    inputType=InputType.Option,
+    inputConfig={
+        "options": ["SQ8", "Flat"],
+    },
+)
+
+CaseConfigParamInput_max_num_levels_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.maxNumLevels,
+    inputType=InputType.Option,
+    inputConfig={
+        "options": [1, 2],
+    },
+)
+
+CaseConfigParamInput_enable_pca_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.enablePca,
+    inputType=InputType.Option,
+    inputConfig={
+        "options": ["on", "off"],
+    },
+)
+
+CaseConfigParamInput_num_leaves_to_search_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.numLeavesToSearch,
+    displayLabel="Num leaves to search",
+    inputHelp="The database flag controls the trade off between recall and QPS",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 20,
+        "max": 10486,
+        "value": 20,
+    },
+)
+
+CaseConfigParamInput_max_top_neighbors_buffer_size_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.maxTopNeighborsBufferSize,
+    displayLabel="Max top neighbors buffer size",
+    inputHelp="The database flag specifies the size of cache used to improve the \
+        performance for filtered queries by scoring or ranking the scanned candidate \
+        neighbors in memory instead of the disk",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 10000,
+        "max": 60000,
+        "value": 20000,
+    },
+)
+
+CaseConfigParamInput_pre_reordering_num_neighbors_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.preReorderingNumNeigbors,
+    displayLabel="Pre reordering num neighbors",
+    inputHelp="Specifies the number of candidate neighbors to consider during the reordering \
+        stages after initial search identifies a set of candidates",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 20,
+        "max": 10486,
+        "value": 80,
+    },
+)
+
+CaseConfigParamInput_num_search_threads_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.numSearchThreads,
+    displayLabel="Num of searcher threads",
+    inputHelp="The number of searcher threads for multi-thread search.",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 1,
+        "max": 100,
+        "value": 2,
+    },
+)
+
+CaseConfigParamInput_max_num_prefetch_datasets_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.maxNumPrefetchDatasets,
+    displayLabel="Max num prefetch datasets",
+    inputHelp="The maximum number of data batches to prefetch during index search, where batch is a group of buffer pages",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 10,
+        "max": 150,
+        "value": 100,
+    },
+)
+
+CaseConfigParamInput_maintenance_work_mem_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.maintenance_work_mem,
+    inputHelp="Recommended value: 1.33x the index size, not to exceed the available free memory."
+    "Specify in gigabytes. e.g. 8GB",
+    inputType=InputType.Text,
+    inputConfig={
+        "value": "8GB",
+    },
+)
+
+CaseConfigParamInput_max_parallel_workers_AlloyDB = CaseConfigInput(
+    label=CaseConfigParamType.max_parallel_workers,
+    displayLabel="Max parallel workers",
+    inputHelp="Recommended value: (cpu cores - 1). This will set the parameters: max_parallel_maintenance_workers,"
+    " max_parallel_workers & table(parallel_workers)",
+    inputType=InputType.Number,
+    inputConfig={
+        "min": 0,
+        "max": 1024,
+        "value": 7,
+    },
+)
+
+
 MilvusLoadConfig = [
     CaseConfigParamInput_IndexType,
     CaseConfigParamInput_M,
@@ -1045,6 +1180,33 @@ PgDiskANNPerformanceConfig = [
     CaseConfigParamInput_l_value_is,
 ]
 
+
+AlloyDBLoadConfig = [
+    CaseConfigParamInput_IndexType_AlloyDB,
+    CaseConfigParamInput_num_leaves_AlloyDB,
+    CaseConfigParamInput_max_num_levels_AlloyDB,
+    CaseConfigParamInput_enable_pca_AlloyDB,
+    CaseConfigParamInput_quantizer_AlloyDB,
+    CaseConfigParamInput_maintenance_work_mem_AlloyDB,
+    CaseConfigParamInput_max_parallel_workers_AlloyDB,
+]
+
+AlloyDBPerformanceConfig = [
+    CaseConfigParamInput_IndexType_AlloyDB,
+    CaseConfigParamInput_num_leaves_AlloyDB,
+    CaseConfigParamInput_max_num_levels_AlloyDB,
+    CaseConfigParamInput_enable_pca_AlloyDB,
+    CaseConfigParamInput_quantizer_AlloyDB,
+    CaseConfigParamInput_num_search_threads_AlloyDB,
+    CaseConfigParamInput_num_leaves_to_search_AlloyDB,
+    CaseConfigParamInput_max_num_prefetch_datasets_AlloyDB,
+    CaseConfigParamInput_max_top_neighbors_buffer_size_AlloyDB,
+    CaseConfigParamInput_pre_reordering_num_neighbors_AlloyDB,
+    CaseConfigParamInput_maintenance_work_mem_AlloyDB,
+    CaseConfigParamInput_max_parallel_workers_AlloyDB,
+]
+
+
 CASE_CONFIG_MAP = {
     DB.Milvus: {
         CaseLabel.Load: MilvusLoadConfig,
@@ -1081,4 +1243,8 @@ CASE_CONFIG_MAP = {
         CaseLabel.Load: PgDiskANNLoadConfig,
         CaseLabel.Performance: PgDiskANNPerformanceConfig,
     },
+    DB.AlloyDB: {
+        CaseLabel.Load: AlloyDBLoadConfig,
+        CaseLabel.Performance: AlloyDBPerformanceConfig,
+    },
 }

vectordb_bench/models.py

@@ -76,6 +76,15 @@ class CaseConfigParamType(Enum):
     num_bits_per_dimension = "num_bits_per_dimension"
     query_search_list_size = "query_search_list_size"
     query_rescore = "query_rescore"
+    numLeaves = "num_leaves"
+    quantizer = "quantizer"
+    enablePca = "enable_pca"
+    maxNumLevels = "max_num_levels"
+    numLeavesToSearch = "num_leaves_to_search"
+    maxTopNeighborsBufferSize = "max_top_neighbors_buffer_size"
+    preReorderingNumNeigbors = "pre_reordering_num_neighbors"
+    numSearchThreads = "num_search_threads"
+    maxNumPrefetchDatasets = "max_num_prefetch_datasets"
 
 
 class CustomizedCase(BaseModel):

{vectordb_bench-0.0.17.dist-info → vectordb_bench-0.0.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectordb-bench
-Version: 0.0.17
+Version: 0.0.18
 Summary: VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze.
 Author-email: XuanYang-cn <xuan.yang@zilliz.com>
 Project-URL: repository, https://github.com/zilliztech/VectorDBBench
@@ -64,6 +64,10 @@ Provides-Extra: pgdiskann
 Requires-Dist: psycopg; extra == "pgdiskann"
 Requires-Dist: psycopg-binary; extra == "pgdiskann"
 Requires-Dist: pgvector; extra == "pgdiskann"
+Provides-Extra: alloydb
+Requires-Dist: psycopg; extra == "alloydb"
+Requires-Dist: psycopg-binary; extra == "alloydb"
+Requires-Dist: pgvector; extra == "alloydb"
 Provides-Extra: pgvecto-rs
 Requires-Dist: pgvecto_rs[psycopg3]>=0.2.2; extra == "pgvecto-rs"
 Provides-Extra: redis
@@ -116,11 +120,12 @@ All the database client supported
 | pgvector                 | `pip install vectordb-bench[pgvector]`      |
 | pgvecto.rs               | `pip install vectordb-bench[pgvecto_rs]`    |
 | pgvectorscale            | `pip install vectordb-bench[pgvectorscale]` |
-| pgdiskann                | `pip install vectordb-bench[pgdiskann]` |
+| pgdiskann                | `pip install vectordb-bench[pgdiskann]`     |
 | redis                    | `pip install vectordb-bench[redis]`         |
 | memorydb                 | `pip install vectordb-bench[memorydb]`      |
 | chromadb                 | `pip install vectordb-bench[chromadb]`      |
 | awsopensearch            | `pip install vectordb-bench[awsopensearch]` |
+| alloydb                  | `pip install vectordb-bench[alloydb]`       |
 
 ### Run
 

{vectordb_bench-0.0.17.dist-info → vectordb_bench-0.0.18.dist-info}/RECORD

@@ -4,7 +4,7 @@ vectordb_bench/base.py,sha256=d34WCGXZI1u5RGQtqrPHd3HbOF5AmioFrM2j30Aj1sY,130
 vectordb_bench/interface.py,sha256=ZT3pseyq--TuxtopdP2hRut-6vIInKo62pvAl2zBD10,9708
 vectordb_bench/log_util.py,sha256=nMnW-sN24WyURcI07t-WA3q2N5R-YIvFgboRsSrNJDg,2906
 vectordb_bench/metric.py,sha256=c-LAxCtb55txnsfd3FN4gRpRif8RREhKRF0eg2JmHGc,2045
-vectordb_bench/models.py,sha256=jAY60QFq3Uhq0YPIQOGYBU1JspOz8D0XJEXSlEBT7cs,10645
+vectordb_bench/models.py,sha256=_0_hscKUqaCHjdjyO_-ntPFgJvgU01y8aldqDcq9ELQ,11041
 vectordb_bench/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectordb_bench/backend/assembler.py,sha256=mmoLzWXFSlrpWvaVY41wiRNWNv2IR-LzlANX55MJbYI,2028
 vectordb_bench/backend/cases.py,sha256=lQ9jgKaJGunj-mJXR3cgGt16wCsrDrvs-GS3ycTDk0U,16169
@@ -13,8 +13,11 @@ vectordb_bench/backend/dataset.py,sha256=MZSu0Q3AkK9gxiuLKNTMH6hhucKK668j4G1-8em
 vectordb_bench/backend/result_collector.py,sha256=jdQf5-q1z5y07SKy9Sig1wFROmm-p9x_Y81fId0sjaU,807
 vectordb_bench/backend/task_runner.py,sha256=cn_RRDyFfNSLlTT84W-ZaXvdl54pK6Cxcsp9ucNRcCs,11864
 vectordb_bench/backend/utils.py,sha256=2UixYyfKvl8zRiashywB1l6hTI3jMtiZhiVm_bXHV1Y,1811
-vectordb_bench/backend/clients/__init__.py,sha256=vl-ldJCTAm6aXUNZS4xE5M2zbiWSz9-rua_9PcOj7ZA,6277
-vectordb_bench/backend/clients/api.py,sha256=KGHdn8gewGm_HbsF7qfn3ibep4AoMXUzaoTjLpF98NE,6176
+vectordb_bench/backend/clients/__init__.py,sha256=JKDRKwfDDZLZFct_wahtSEpqu00HVZBNsVyhJ_hHTL8,6701
+vectordb_bench/backend/clients/api.py,sha256=phvqTME3NEPyZGTo85MPeOWwICZO06W3388GT5g72pc,6210
+vectordb_bench/backend/clients/alloydb/alloydb.py,sha256=rAV558tyd2hX3jcl3bRcxOkeq__GSAXLxfl3MqkAVkM,13375
+vectordb_bench/backend/clients/alloydb/cli.py,sha256=5g3heAEfuwIHCUIHDU4LYQq-CaQto0sGAdr45jdhuNc,4970
+vectordb_bench/backend/clients/alloydb/config.py,sha256=JFQMHvBWG1P5T4N7B95o4tMfN4cVqb01I5TNvjDYQuw,5358
 vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py,sha256=O42OU7K7L0KcJ96AphjmgkyN7a220VLsOL0QwghY2aw,8038
 vectordb_bench/backend/clients/aws_opensearch/cli.py,sha256=v1bGoovgokhIGN5tZwb_MrP4af7BfXYQaOpDuy0Ibh0,1327
 vectordb_bench/backend/clients/aws_opensearch/config.py,sha256=USytbPX0zDjcWTg51J37vWWYsgualPcCHGfRFdzmGtg,1890
@@ -28,7 +31,7 @@ vectordb_bench/backend/clients/memorydb/config.py,sha256=PjhLMMr_LdJ8O91JpHNCCT6
 vectordb_bench/backend/clients/memorydb/memorydb.py,sha256=XIqtXpY-2lJohIuImFDsRO3c_upn04eCplIOlaLxFo4,10114
 vectordb_bench/backend/clients/milvus/cli.py,sha256=QqzYIOeUSXEvdLH0_YUMhwDHUDJirTNKeUxrJQIqSdw,8506
 vectordb_bench/backend/clients/milvus/config.py,sha256=AZ4QHoufRIjsX2eVrtnug8SeYnuHeBMna_34OQNFxz0,6847
-vectordb_bench/backend/clients/milvus/milvus.py,sha256=BhEkJr8ZQuiFqYd1sQYhKd8YXHS9vlaqOv36zlHI6xc,7712
+vectordb_bench/backend/clients/milvus/milvus.py,sha256=7l2ilpZeCVnXLkWjut2EoIhUraYL7qWfSsBrRaq4vuo,7700
 vectordb_bench/backend/clients/pgdiskann/cli.py,sha256=ued1DyufltataIk6KcmBkNp8PdB9Aj65nVJ6WhrD_VI,3130
 vectordb_bench/backend/clients/pgdiskann/config.py,sha256=8E0GLgUxa5LlJ_eXCugbbO08qdbCVqc1wtdsoOsKEW4,4444
 vectordb_bench/backend/clients/pgdiskann/pgdiskann.py,sha256=bEcbpTVSFxRJ5HiJTX77cgu6NqTMPs8qiGeMF7jBC30,12628
@@ -65,7 +68,7 @@ vectordb_bench/backend/runner/serial_runner.py,sha256=ku1Dtps9JcmwCwZq7eDw0pcP9I
 vectordb_bench/backend/runner/util.py,sha256=pGJn-qXWwGXVlmsMulaqH0zXcasDWjsVwwOJeDFWXhc,1032
 vectordb_bench/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectordb_bench/cli/cli.py,sha256=Z2-vLwvnnZFsVAPyjFK557cZZYWX_q60XVJP-aYUGdc,15416
-vectordb_bench/cli/vectordbbench.py,sha256=QOKwURuiaxsQ-vf7ji-qHEug5j1y_j1xSeykSsENz30,1055
+vectordb_bench/cli/vectordbbench.py,sha256=jbpyjh4xKVRocxg4XurLL3ABUzBRXEChRGYhyqH4ItE,1140
 vectordb_bench/config-files/sample_config.yml,sha256=yw9ZgHczNi9PedNuTVxZKiOTI6AVoQS1h8INNgoDjPk,340
 vectordb_bench/custom/custom_case.json,sha256=uKo7NJgXDPPLtf_V6y1uc5w1aIcjLp-GCJEYOCty1As,475
 vectordb_bench/frontend/utils.py,sha256=jCyfk0QyLl3RLh-1MBRbBy6aep9hO32ScJCDYA2kaZU,489
@@ -94,7 +97,7 @@ vectordb_bench/frontend/components/run_test/hideSidebar.py,sha256=vb5kzIMmbMqWX6
 vectordb_bench/frontend/components/run_test/initStyle.py,sha256=osPUgfFfH7rRlVNHSMumvmZxvKWlLxmZiNqgnMiUJEU,723
 vectordb_bench/frontend/components/run_test/submitTask.py,sha256=NCEXfR3xudAncjVEvsV2iaiov5AatGObe830UI6481M,3341
 vectordb_bench/frontend/components/tables/data.py,sha256=pVG_hb4bTMLfUt10NUCJSqcFkPmnN7i9jTw9DcWizpI,1364
-vectordb_bench/frontend/config/dbCaseConfigs.py,sha256=Sl6sPTtUJki8uRu5wgPynKcR4OdaDwIgnmzCVVX9gQ0,31070
+vectordb_bench/frontend/config/dbCaseConfigs.py,sha256=40rJ3YtRkPbZBBlWnV7csXDArqrMVaBw6tLMtxzj3w4,36311
 vectordb_bench/frontend/config/dbPrices.py,sha256=10aBKjVcEg8y7TPSda28opmBM1KmXNrvbU9WM_BsZcE,176
 vectordb_bench/frontend/config/styles.py,sha256=E2PmwmiewxBKJJ59hQ4ZXatqg8QTN-Z53JlsvWMHM2M,2291
 vectordb_bench/frontend/pages/concurrent.py,sha256=z2izkQ0suO5mZ8PpVY2jypZkF5VT8xUkQQEkwd6C-ww,2094
@@ -120,9 +123,9 @@ vectordb_bench/results/WeaviateCloud/result_20230808_standard_weaviatecloud.json
 vectordb_bench/results/ZillizCloud/result_20230727_standard_zillizcloud.json,sha256=-Mdm4By65XDRCrmVOCF8yQXjcZtH4Xo4shcjoDoBUKU,18293
 vectordb_bench/results/ZillizCloud/result_20230808_standard_zillizcloud.json,sha256=77XlHT5zM_K7mG5HfDQKwXZnSCuR37VUbt6-P3J_amI,15737
 vectordb_bench/results/ZillizCloud/result_20240105_standard_202401_zillizcloud.json,sha256=TualfJ0664Hs-vdIW68bdkqAEYyzotXmu2P0yIN-GHk,42526
-vectordb_bench-0.0.17.dist-info/LICENSE,sha256=HXbxhrb5u5SegVzeLNF_voVgRsJMavcLaOmD1N0lZkM,1067
-vectordb_bench-0.0.17.dist-info/METADATA,sha256=SreMZtCcdr-dKZrwyLXuVHX-WinJe9uLTvV9_NtIXgY,34643
-vectordb_bench-0.0.17.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-vectordb_bench-0.0.17.dist-info/entry_points.txt,sha256=Qzw6gVx96ui8esG21H6yHsI6nboEohRmV424TYhQNrA,113
-vectordb_bench-0.0.17.dist-info/top_level.txt,sha256=jnhZFZAuKX1J60yt-XOeBZ__ctiZMvoC_s0RFq29lpM,15
-vectordb_bench-0.0.17.dist-info/RECORD,,
+vectordb_bench-0.0.18.dist-info/LICENSE,sha256=HXbxhrb5u5SegVzeLNF_voVgRsJMavcLaOmD1N0lZkM,1067
+vectordb_bench-0.0.18.dist-info/METADATA,sha256=PtTRr7G1PXhzA88gZYML3Y6JPaqzXk_ejeFbvoRVMOQ,34883
+vectordb_bench-0.0.18.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+vectordb_bench-0.0.18.dist-info/entry_points.txt,sha256=Qzw6gVx96ui8esG21H6yHsI6nboEohRmV424TYhQNrA,113
+vectordb_bench-0.0.18.dist-info/top_level.txt,sha256=jnhZFZAuKX1J60yt-XOeBZ__ctiZMvoC_s0RFq29lpM,15
+vectordb_bench-0.0.18.dist-info/RECORD,,