vectordb-bench 0.0.16__tar.gz → 0.0.18__tar.gz

{vectordb_bench-0.0.16 → vectordb_bench-0.0.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectordb-bench
-Version: 0.0.16
+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.16 → vectordb_bench-0.0.18}/README.md

@@ -38,11 +38,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.16 → vectordb_bench-0.0.18}/pyproject.toml

@@ -73,6 +73,7 @@ elastic = [ "elasticsearch" ]
 pgvector = [ "psycopg", "psycopg-binary", "pgvector" ]
 pgvectorscale = [ "psycopg", "psycopg-binary", "pgvector" ]
 pgdiskann = [ "psycopg", "psycopg-binary", "pgvector" ]
+alloydb = [ "psycopg", "psycopg-binary", "pgvector"]
 pgvecto_rs = [ "pgvecto_rs[psycopg3]>=0.2.2" ]
 redis = [ "redis" ]
 memorydb = [ "memorydb" ]

vectordb_bench-0.0.18/tests/test_rate_runner.py

@@ -0,0 +1,88 @@
+from typing import Iterable
+import argparse
+from vectordb_bench.backend.dataset import Dataset, DatasetSource
+from vectordb_bench.backend.runner.rate_runner import RatedMultiThreadingInsertRunner
+from vectordb_bench.backend.runner.read_write_runner import ReadWriteRunner
+from vectordb_bench.backend.clients import DB, VectorDB
+from vectordb_bench.backend.clients.milvus.config import FLATConfig
+from vectordb_bench.backend.clients.zilliz_cloud.config import AutoIndexConfig
+
+import logging
+
+log = logging.getLogger("vectordb_bench")
+log.setLevel(logging.DEBUG)
+
+def get_rate_runner(db):
+    cohere = Dataset.COHERE.manager(100_000)
+    prepared = cohere.prepare(DatasetSource.AliyunOSS)
+    assert prepared
+    runner = RatedMultiThreadingInsertRunner(
+        rate = 10,
+        db = db,
+        dataset = cohere,
+    )
+
+    return runner
+
+def test_rate_runner(db, insert_rate):
+    runner = get_rate_runner(db)
+
+    _, t = runner.run_with_rate()
+    log.info(f"insert run done, time={t}")
+
+def test_read_write_runner(db, insert_rate, conc: list, search_stage: Iterable[float], read_dur_after_write: int, local: bool=False):
+    cohere = Dataset.COHERE.manager(1_000_000)
+    if local is True:
+        source = DatasetSource.AliyunOSS
+    else:
+        source = DatasetSource.S3
+    prepared = cohere.prepare(source)
+    assert prepared
+
+    rw_runner = ReadWriteRunner(
+        db=db,
+        dataset=cohere,
+        insert_rate=insert_rate,
+        search_stage=search_stage,
+        read_dur_after_write=read_dur_after_write,
+        concurrencies=conc
+    )
+    rw_runner.run_read_write()
+
+
+def get_db(db: str, config: dict) -> VectorDB:
+    if db == DB.Milvus.name:
+        return DB.Milvus.init_cls(dim=768, db_config=config, db_case_config=FLATConfig(metric_type="COSINE"), drop_old=True, pre_load=True)
+    elif db == DB.ZillizCloud.name:
+        return DB.ZillizCloud.init_cls(dim=768, db_config=config, db_case_config=AutoIndexConfig(metric_type="COSINE"), drop_old=True, pre_load=True)
+    else:
+        raise ValueError(f"unknown db: {db}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-r", "--insert_rate", type=int, default="1000", help="insert entity row count per seconds, cps")
+    parser.add_argument("-d", "--db", type=str, default=DB.Milvus.name, help="db name")
+    parser.add_argument("-t", "--duration", type=int, default=300, help="stage search duration in seconds")
+    parser.add_argument("--use_s3", action='store_true', help="whether to use S3 dataset")
+
+    flags = parser.parse_args()
+
+    # TODO read uri, user, password from .env
+    config = {
+        "uri": "http://localhost:19530",
+        "user": "",
+        "password": "",
+    }
+
+    conc = (1, 15, 50)
+    search_stage = (0.5, 0.6, 0.7, 0.8, 0.9, 1.0)
+
+    db = get_db(flags.db, config)
+    test_read_write_runner(
+        db=db,
+        insert_rate=flags.insert_rate,
+        conc=conc,
+        search_stage=search_stage,
+        read_dur_after_write=flags.duration,
+        local=flags.use_s3)

{vectordb_bench-0.0.16 → vectordb_bench-0.0.18}/vectordb_bench/__init__.py

@@ -17,7 +17,7 @@ class config:
 
     DEFAULT_DATASET_URL = env.str("DEFAULT_DATASET_URL", AWS_S3_URL)
     DATASET_LOCAL_DIR = env.path("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset")
-    NUM_PER_BATCH = env.int("NUM_PER_BATCH", 5000)
+    NUM_PER_BATCH = env.int("NUM_PER_BATCH", 100)
 
     DROP_OLD = env.bool("DROP_OLD", True)
     USE_SHUFFLED_DATA = env.bool("USE_SHUFFLED_DATA", True)

{vectordb_bench-0.0.16 → vectordb_bench-0.0.18}/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-0.0.18/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()]