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

vectordb_bench/backend/clients/aws_opensearch/run.py

@@ -0,0 +1,125 @@
+import time, random
+from opensearchpy import OpenSearch
+from opensearch_dsl import Search, Document, Text, Keyword
+
+_HOST = 'xxxxxx.us-west-2.es.amazonaws.com'
+_PORT = 443
+_AUTH = ('admin', 'xxxxxx') # For testing only. Don't store credentials in code.
+
+_INDEX_NAME = 'my-dsl-index'
+_BATCH = 100
+_ROWS = 100
+_DIM = 128
+_TOPK = 10
+
+
+def create_client():
+    client = OpenSearch(
+        hosts=[{'host': _HOST, 'port': _PORT}],
+        http_compress=True, # enables gzip compression for request bodies
+        http_auth=_AUTH,
+        use_ssl=True,
+        verify_certs=True,
+        ssl_assert_hostname=False,
+        ssl_show_warn=False,
+    )
+    return client
+
+
+def create_index(client, index_name):
+    settings = {
+        "index": {
+            "knn": True,
+            "number_of_shards": 1,
+            "refresh_interval": "5s",
+        }
+    }
+    mappings = {
+        "properties": {
+            "embedding": {
+                "type": "knn_vector",
+                "dimension": _DIM,
+                "method": {
+                    "engine": "nmslib",
+                    "name": "hnsw",
+                    "space_type": "l2",
+                    "parameters": {
+                        "ef_construction": 128,
+                        "m": 24,
+                    }
+                }
+            }
+        }
+    }
+
+    response = client.indices.create(index=index_name, body=dict(settings=settings, mappings=mappings))
+    print('\nCreating index:')
+    print(response)
+
+
+def delete_index(client, index_name):
+    response = client.indices.delete(index=index_name)
+    print('\nDeleting index:')
+    print(response)
+
+
+def bulk_insert(client, index_name):
+    # Perform bulk operations
+    ids = [i for i in range(_ROWS)]
+    vec = [[random.random() for _ in range(_DIM)] for _ in range(_ROWS)]
+
+    docs = []
+    for i in range(0, _ROWS, _BATCH):
+        docs.clear()
+        for j in range(0, _BATCH):
+            docs.append({"index": {"_index": index_name, "_id": ids[i+j]}})
+            docs.append({"embedding": vec[i+j]})
+        response = client.bulk(docs)
+        print('\nAdding documents:', len(response['items']), response['errors'])
+        response = client.indices.stats(index_name)
+        print('\nTotal document count in index:', response['_all']['primaries']['indexing']['index_total'])
+
+
+def search(client, index_name):
+    # Search for the document.
+    search_body = {
+        "size": _TOPK,
+        "query": {
+            "knn": {
+                "embedding": {
+                    "vector": [random.random() for _ in range(_DIM)],
+                    "k": _TOPK,
+                }
+            }
+        }
+    }
+    while True:
+        response = client.search(index=index_name, body=search_body)
+        print(f'\nSearch took: {response["took"]}')
+        print(f'\nSearch shards: {response["_shards"]}')
+        print(f'\nSearch hits total: {response["hits"]["total"]}')
+        result = response["hits"]["hits"]
+        if len(result) != 0:
+            print('\nSearch results:')
+            for hit in response["hits"]["hits"]:
+                print(hit["_id"], hit["_score"])
+            break
+        else:
+            print('\nSearch not ready, sleep 1s')
+            time.sleep(1)
+
+
+def main():
+    client = create_client()
+    try:
+        create_index(client, _INDEX_NAME)
+        bulk_insert(client, _INDEX_NAME)
+        search(client, _INDEX_NAME)
+        delete_index(client, _INDEX_NAME)
+    except Exception as e:
+        print(e)
+        delete_index(client, _INDEX_NAME)
+
+
+if __name__ == '__main__':
+    main()

vectordb_bench/backend/clients/milvus/cli.py

@@ -0,0 +1,291 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor3,
+    IVFFlatTypedDictN,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+
+)
+from vectordb_bench.backend.clients import DB
+
+DBTYPE = DB.Milvus
+
+
+class MilvusTypedDict(TypedDict):
+    uri: Annotated[
+        str, click.option("--uri", type=str, help="uri connection string", required=True)
+    ]
+
+
+class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
+def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]):
+    from .config import MilvusConfig, AutoIndexConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=AutoIndexConfig(),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict)
+def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]):
+    from .config import MilvusConfig, FLATConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=FLATConfig(),
+        **parameters,
+    )
+
+
+class MilvusHNSWTypedDict(CommonTypedDict, MilvusTypedDict, HNSWFlavor3):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusHNSWTypedDict)
+def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]):
+    from .config import MilvusConfig, HNSWConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=HNSWConfig(
+            M=parameters["m"],
+            efConstruction=parameters["ef_construction"],
+            ef=parameters["ef_search"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusIVFFlatTypedDict(CommonTypedDict, MilvusTypedDict, IVFFlatTypedDictN):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
+def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]):
+    from .config import MilvusConfig, IVFFlatConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=IVFFlatConfig(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+        ),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusIVFFlatTypedDict)
+def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]):
+    from .config import MilvusConfig, IVFSQ8Config
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=IVFSQ8Config(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusDISKANNTypedDict(CommonTypedDict, MilvusTypedDict):
+    search_list: Annotated[
+        str, click.option("--search-list",
+                          type=int,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusDISKANNTypedDict)
+def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]):
+    from .config import MilvusConfig, DISKANNConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=DISKANNConfig(
+            search_list=parameters["search_list"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusGPUIVFTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict):
+    cache_dataset_on_device: Annotated[
+        str, click.option("--cache-dataset-on-device",
+                          type=str,
+                          required=True)
+    ]
+    refine_ratio: Annotated[
+        str, click.option("--refine-ratio",
+                          type=float,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUIVFTypedDict)
+def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]):
+    from .config import MilvusConfig, GPUIVFFlatConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=GPUIVFFlatConfig(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+            cache_dataset_on_device=parameters["cache_dataset_on_device"],
+            refine_ratio=parameters.get("refine_ratio"),
+        ),
+        **parameters,
+    )
+
+
+class MilvusGPUIVFPQTypedDict(CommonTypedDict, MilvusTypedDict, MilvusIVFFlatTypedDict, MilvusGPUIVFTypedDict):
+    m: Annotated[
+        str, click.option("--m",
+                          type=int, help="hnsw m",
+                          required=True)
+    ]
+    nbits: Annotated[
+        str, click.option("--nbits",
+                          type=int,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUIVFPQTypedDict)
+def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]):
+    from .config import MilvusConfig, GPUIVFPQConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=GPUIVFPQConfig(
+            nlist=parameters["nlist"],
+            nprobe=parameters["nprobe"],
+            m=parameters["m"],
+            nbits=parameters["nbits"],
+            cache_dataset_on_device=parameters["cache_dataset_on_device"],
+            refine_ratio=parameters["refine_ratio"],
+        ),
+        **parameters,
+    )
+
+
+class MilvusGPUCAGRATypedDict(CommonTypedDict, MilvusTypedDict, MilvusGPUIVFTypedDict):
+    intermediate_graph_degree: Annotated[
+        str, click.option("--intermediate-graph-degree",
+                          type=int,
+                          required=True)
+    ]
+    graph_degree: Annotated[
+        str, click.option("--graph-degree",
+                          type=int,
+                          required=True)
+    ]
+    build_algo: Annotated[
+        str, click.option("--build_algo",
+                          type=str,
+                          required=True)
+    ]
+    team_size: Annotated[
+        str, click.option("--team-size",
+                          type=int,
+                          required=True)
+    ]
+    search_width: Annotated[
+        str, click.option("--search-width",
+                          type=int,
+                          required=True)
+    ]
+    itopk_size: Annotated[
+        str, click.option("--itopk-size",
+                          type=int,
+                          required=True)
+    ]
+    min_iterations: Annotated[
+        str, click.option("--min-iterations",
+                          type=int,
+                          required=True)
+    ]
+    max_iterations: Annotated[
+        str, click.option("--max-iterations",
+                          type=int,
+                          required=True)
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(MilvusGPUCAGRATypedDict)
+def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]):
+    from .config import MilvusConfig, GPUCAGRAConfig
+
+    run(
+        db=DBTYPE,
+        db_config=MilvusConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+        ),
+        db_case_config=GPUCAGRAConfig(
+            intermediate_graph_degree=parameters["intermediate_graph_degree"],
+            graph_degree=parameters["graph_degree"],
+            itopk_size=parameters["itopk_size"],
+            team_size=parameters["team_size"],
+            search_width=parameters["search_width"],
+            min_iterations=parameters["min_iterations"],
+            max_iterations=parameters["max_iterations"],
+            build_algo=parameters["build_algo"],
+            cache_dataset_on_device=parameters["cache_dataset_on_device"],
+            refine_ratio=parameters["refine_ratio"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/milvus/milvus.py

@@ -8,7 +8,7 @@ from typing import Iterable
 from pymilvus import Collection, utility
 from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusException
 
-from ..api import VectorDB
+from ..api import VectorDB, IndexType
 from .config import MilvusIndexConfig
 
 
@@ -122,10 +122,18 @@ class Milvus(VectorDB):
             if self.case_config.is_gpu_index:
                 log.debug("skip compaction for gpu index type.")
             else :
-                self.col.compact()
-                self.col.wait_for_compaction_completed()
+                try:
+                    self.col.compact()
+                    self.col.wait_for_compaction_completed()
+                except Exception as e:
+                    log.warning(f"{self.name} compact error: {e}")
+                    if hasattr(e, 'code'):
+                        if e.code().name == 'PERMISSION_DENIED':
+                            log.warning(f"Skip compact due to permission denied.")
+                            pass
+                    else:
+                        raise e
                 wait_index()
-
         except Exception as e:
             log.warning(f"{self.name} optimize error: {e}")
             raise e from None
@@ -143,7 +151,6 @@ class Milvus(VectorDB):
                     self.case_config.index_param(),
                     index_name=self._index_name,
                 )
-
             coll.load()
             log.info(f"{self.name} load")
         except Exception as e:
@@ -160,7 +167,7 @@ class Milvus(VectorDB):
         if self.case_config.is_gpu_index:
             log.info(f"current gpu_index only supports IP / L2, cosine dataset need normalize.")
             return True
-        
+
         return False
 
     def insert_embeddings(

vectordb_bench/backend/clients/pgvector/cli.py

@@ -0,0 +1,116 @@
+from typing import Annotated, Optional, TypedDict, Unpack
+
+import click
+import os
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor1,
+    IVFFlatTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from vectordb_bench.backend.clients import DB
+
+
+class PgVectorTypedDict(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 PgVectorIVFFlatTypedDict(PgVectorTypedDict, IVFFlatTypedDict):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectorIVFFlatTypedDict)
+def PgVectorIVFFlat(
+    **parameters: Unpack[PgVectorIVFFlatTypedDict],
+):
+    from .config import PgVectorConfig, PgVectorIVFFlatConfig
+
+    run(
+        db=DB.PgVector,
+        db_config=PgVectorConfig(
+            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=PgVectorIVFFlatConfig(
+            metric_type=None, lists=parameters["lists"], probes=parameters["probes"]
+        ),
+        **parameters,
+    )
+
+
+class PgVectorHNSWTypedDict(PgVectorTypedDict, HNSWFlavor1):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgVectorHNSWTypedDict)
+def PgVectorHNSW(
+    **parameters: Unpack[PgVectorHNSWTypedDict],
+):
+    from .config import PgVectorConfig, PgVectorHNSWConfig
+
+    run(
+        db=DB.PgVector,
+        db_config=PgVectorConfig(
+            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=PgVectorHNSWConfig(
+            m=parameters["m"],
+            ef_construction=parameters["ef_construction"],
+            ef_search=parameters["ef_search"],
+            maintenance_work_mem=parameters["maintenance_work_mem"],
+            max_parallel_workers=parameters["max_parallel_workers"],
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/pgvector/config.py

@@ -109,7 +109,7 @@ class PgVectorIndexConfig(BaseModel, DBCaseConfig):
     def _optionally_build_set_options(
         set_mapping: Mapping[str, Any]
     ) -> Sequence[dict[str, Any]]:
-        """Walk through options, creating 'SET 'key1 = "value1";' commands"""
+        """Walk through options, creating 'SET 'key1 = "value1";' list"""
         session_options = []
         for setting_name, value in set_mapping.items():
             if value:

vectordb_bench/backend/clients/pgvector/pgvector.py

@@ -58,14 +58,13 @@ class PgVector(VectorDB):
                 self.case_config.create_index_after_load,
             )
         ):
-            err = f"{self.name} config must create an index using create_index_before_load and/or 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.pg_table.drop(pg_engine, checkfirst=True)
             self._drop_index()
             self._drop_table()
             self._create_table(dim)
@@ -257,7 +256,10 @@ class PgVector(VectorDB):
             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})"
+            """
+            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),
@@ -339,9 +341,10 @@ class PgVector(VectorDB):
         assert self.conn is not None, "Connection is not initialized"
         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, (query, k), prepare=True, binary=True
+            self._unfiltered_search, (q, k), prepare=True, binary=True
         )
 
         return [int(i[0]) for i in result.fetchall()]

vectordb_bench/backend/clients/redis/cli.py

@@ -0,0 +1,74 @@
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor2,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+
+
+class RedisTypedDict(TypedDict):
+    host: Annotated[
+        str, click.option("--host", type=str, help="Db host", required=True)
+    ]
+    password: Annotated[str, click.option("--password", type=str, help="Db password")]
+    port: Annotated[int, click.option("--port", type=int, default=6379, help="Db Port")]
+    ssl: Annotated[
+        bool,
+        click.option(
+            "--ssl/--no-ssl",
+            is_flag=True,
+            show_default=True,
+            default=True,
+            help="Enable or disable SSL for Redis",
+        ),
+    ]
+    ssl_ca_certs: Annotated[
+        str,
+        click.option(
+            "--ssl-ca-certs",
+            show_default=True,
+            help="Path to certificate authority file to use for SSL",
+        ),
+    ]
+    cmd: Annotated[
+        bool,
+        click.option(
+            "--cmd",
+            is_flag=True,
+            show_default=True,
+            default=False,
+            help="Cluster Mode Disabled (CMD) for Redis doesn't use Cluster conn",
+        ),
+    ]
+
+
+class RedisHNSWTypedDict(CommonTypedDict, RedisTypedDict, HNSWFlavor2):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(RedisHNSWTypedDict)
+def Redis(**parameters: Unpack[RedisHNSWTypedDict]):
+    from .config import RedisConfig
+    run(
+        db=DB.Redis,
+        db_config=RedisConfig(
+            db_label=parameters["db_label"],
+            password=SecretStr(parameters["password"])
+            if parameters["password"]
+            else None,
+            host=SecretStr(parameters["host"]),
+            port=parameters["port"],
+            ssl=parameters["ssl"],
+            ssl_ca_certs=parameters["ssl_ca_certs"],
+            cmd=parameters["cmd"],
+        ),
+        **parameters,
+    )