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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectordb-bench
-Version: 0.0.14
+Version: 0.0.16
 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
@@ -60,6 +60,10 @@ Provides-Extra: pgvectorscale
 Requires-Dist: psycopg; extra == "pgvectorscale"
 Requires-Dist: psycopg-binary; extra == "pgvectorscale"
 Requires-Dist: pgvector; extra == "pgvectorscale"
+Provides-Extra: pgdiskann
+Requires-Dist: psycopg; extra == "pgdiskann"
+Requires-Dist: psycopg-binary; extra == "pgdiskann"
+Requires-Dist: pgvector; extra == "pgdiskann"
 Provides-Extra: pgvecto-rs
 Requires-Dist: pgvecto_rs[psycopg3]>=0.2.2; extra == "pgvecto-rs"
 Provides-Extra: redis
@@ -112,6 +116,7 @@ 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]` |
 | redis                    | `pip install vectordb-bench[redis]`         |
 | memorydb                 | `pip install vectordb-bench[memorydb]`      |
 | chromadb                 | `pip install vectordb-bench[chromadb]`      |
@@ -191,7 +196,7 @@ Options:
   --m INTEGER                     hnsw m
   --ef-construction INTEGER       hnsw ef-construction
   --ef-search INTEGER             hnsw ef-search
-  --quantization-type [none|halfvec]
+  --quantization-type [none|bit|halfvec]
                                   quantization type for vectors
   --custom-case-name TEXT         Custom case name i.e. PerformanceCase1536D50K
   --custom-case-description TEXT  Custom name description

{vectordb_bench-0.0.14 → vectordb_bench-0.0.16}/README.md

@@ -38,6 +38,7 @@ 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]` |
 | redis                    | `pip install vectordb-bench[redis]`         |
 | memorydb                 | `pip install vectordb-bench[memorydb]`      |
 | chromadb                 | `pip install vectordb-bench[chromadb]`      |
@@ -117,7 +118,7 @@ Options:
   --m INTEGER                     hnsw m
   --ef-construction INTEGER       hnsw ef-construction
   --ef-search INTEGER             hnsw ef-search
-  --quantization-type [none|halfvec]
+  --quantization-type [none|bit|halfvec]
                                   quantization type for vectors
   --custom-case-name TEXT         Custom case name i.e. PerformanceCase1536D50K
   --custom-case-description TEXT  Custom name description

{vectordb_bench-0.0.14 → vectordb_bench-0.0.16}/pyproject.toml

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

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

@@ -37,23 +37,24 @@ class config:
     K_DEFAULT = 100  # default return top k nearest neighbors during search
     CUSTOM_CONFIG_DIR = pathlib.Path(__file__).parent.joinpath("custom/custom_case.json")
 
-    CAPACITY_TIMEOUT_IN_SECONDS = 24 * 3600 # 24h
-    LOAD_TIMEOUT_DEFAULT        = 2.5 * 3600 # 2.5h
-    LOAD_TIMEOUT_768D_1M        = 2.5 * 3600 # 2.5h
-    LOAD_TIMEOUT_768D_10M       =  25 * 3600 # 25h
-    LOAD_TIMEOUT_768D_100M      = 250 * 3600 # 10.41d
+    CAPACITY_TIMEOUT_IN_SECONDS = 24 * 3600   # 24h
+    LOAD_TIMEOUT_DEFAULT        = 24 * 3600   # 24h
+    LOAD_TIMEOUT_768D_1M        = 24 * 3600   # 24h
+    LOAD_TIMEOUT_768D_10M       = 240 * 3600  # 10d
+    LOAD_TIMEOUT_768D_100M      = 2400 * 3600 # 100d
 
-    LOAD_TIMEOUT_1536D_500K     = 2.5 * 3600 # 2.5h
-    LOAD_TIMEOUT_1536D_5M       =  25 * 3600 # 25h
+    LOAD_TIMEOUT_1536D_500K     = 24 * 3600   # 24h
+    LOAD_TIMEOUT_1536D_5M       = 240 * 3600  # 10d
 
-    OPTIMIZE_TIMEOUT_DEFAULT    = 30 * 60   # 30min
-    OPTIMIZE_TIMEOUT_768D_1M    =  30 * 60   # 30min
-    OPTIMIZE_TIMEOUT_768D_10M   = 5 * 3600 # 5h
-    OPTIMIZE_TIMEOUT_768D_100M  =  50 * 3600 # 50h
+    OPTIMIZE_TIMEOUT_DEFAULT    = 24 * 3600   # 24h
+    OPTIMIZE_TIMEOUT_768D_1M    = 24 * 3600   # 24h
+    OPTIMIZE_TIMEOUT_768D_10M   = 240 * 3600  # 10d
+    OPTIMIZE_TIMEOUT_768D_100M  = 2400 * 3600 # 100d
 
 
-    OPTIMIZE_TIMEOUT_1536D_500K =  15 * 60   # 15min
-    OPTIMIZE_TIMEOUT_1536D_5M   =   2.5 * 3600 # 2.5h
+    OPTIMIZE_TIMEOUT_1536D_500K = 24 * 3600   # 24h
+    OPTIMIZE_TIMEOUT_1536D_5M   = 240 * 3600  # 10d
+    
     def display(self) -> str:
         tmp = [
             i for i in inspect.getmembers(self)

{vectordb_bench-0.0.14 → vectordb_bench-0.0.16}/vectordb_bench/backend/clients/__init__.py

@@ -31,6 +31,7 @@ class DB(Enum):
     PgVector = "PgVector"
     PgVectoRS = "PgVectoRS"
     PgVectorScale = "PgVectorScale"
+    PgDiskANN = "PgDiskANN"
     Redis = "Redis"
     MemoryDB = "MemoryDB"
     Chroma = "Chroma"
@@ -77,6 +78,10 @@ class DB(Enum):
             from .pgvectorscale.pgvectorscale import PgVectorScale
             return PgVectorScale
 
+        if self == DB.PgDiskANN:
+            from .pgdiskann.pgdiskann import PgDiskANN
+            return PgDiskANN
+
         if self == DB.Redis:
             from .redis.redis import Redis
             return Redis
@@ -132,6 +137,10 @@ class DB(Enum):
             from .pgvectorscale.config import PgVectorScaleConfig
             return PgVectorScaleConfig
 
+        if self == DB.PgDiskANN:
+            from .pgdiskann.config import PgDiskANNConfig
+            return PgDiskANNConfig
+
         if self == DB.Redis:
             from .redis.config import RedisConfig
             return RedisConfig
@@ -185,6 +194,10 @@ class DB(Enum):
             from .pgvectorscale.config import _pgvectorscale_case_config
             return _pgvectorscale_case_config.get(index_type)
 
+        if self == DB.PgDiskANN:
+            from .pgdiskann.config import _pgdiskann_case_config
+            return _pgdiskann_case_config.get(index_type)
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

{vectordb_bench-0.0.14 → vectordb_bench-0.0.16}/vectordb_bench/backend/clients/api.py

@@ -10,6 +10,8 @@ class MetricType(str, Enum):
     L2 = "L2"
     COSINE = "COSINE"
     IP = "IP"
+    HAMMING = "HAMMING"
+    JACCARD = "JACCARD"
 
 
 class IndexType(str, Enum):

vectordb_bench-0.0.16/vectordb_bench/backend/clients/pgdiskann/cli.py

@@ -0,0 +1,99 @@
+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, Optional, Unpack
+from vectordb_bench.backend.clients import DB
+
+
+class PgDiskAnnTypedDict(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)
+    ]
+    max_neighbors: Annotated[
+        int,
+        click.option(
+            "--max-neighbors", type=int, help="PgDiskAnn max neighbors",
+        ),
+    ]
+    l_value_ib: Annotated[
+        int,
+        click.option(
+            "--l-value-ib", type=int, help="PgDiskAnn l_value_ib",
+        ),
+    ]
+    l_value_is: Annotated[
+        float,
+        click.option(
+            "--l-value-is", type=float, help="PgDiskAnn l_value_is",
+        ),
+    ]
+    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,
+        ),
+    ]
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(PgDiskAnnTypedDict)
+def PgDiskAnn(
+    **parameters: Unpack[PgDiskAnnTypedDict],
+):
+    from .config import PgDiskANNConfig, PgDiskANNImplConfig
+
+    run(
+        db=DB.PgDiskANN,
+        db_config=PgDiskANNConfig(
+            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=PgDiskANNImplConfig(
+            max_neighbors=parameters["max_neighbors"],
+            l_value_ib=parameters["l_value_ib"],
+            l_value_is=parameters["l_value_is"],
+            max_parallel_workers=parameters["max_parallel_workers"],
+            maintenance_work_mem=parameters["maintenance_work_mem"],
+        ),
+        **parameters,
+    )

vectordb_bench-0.0.16/vectordb_bench/backend/clients/pgdiskann/config.py

@@ -0,0 +1,145 @@
+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 PgDiskANNConfigDict(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 PgDiskANNConfig(DBConfig):
+    user_name: SecretStr = SecretStr("postgres")
+    password: SecretStr
+    host: str = "localhost"
+    port: int = 5432
+    db_name: str
+
+    def to_dict(self) -> PgDiskANNConfigDict:
+        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 PgDiskANNIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    create_index_before_load: bool = False
+    create_index_after_load: bool = True
+    maintenance_work_mem: Optional[str]
+    max_parallel_workers: Optional[int]
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "vector_l2_ops"
+        elif self.metric_type == MetricType.IP:
+            return "vector_ip_ops"
+        return "vector_cosine_ops"
+
+    def parse_metric_fun_op(self) -> LiteralString:
+        if self.metric_type == MetricType.L2:
+            return "<->"
+        elif self.metric_type == MetricType.IP:
+            return "<#>"
+        return "<=>"
+
+    def parse_metric_fun_str(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "l2_distance"
+        elif self.metric_type == MetricType.IP:
+            return "max_inner_product"
+        return "cosine_distance"
+    
+    @abstractmethod
+    def index_param(self) -> dict:
+        ...
+
+    @abstractmethod
+    def search_param(self) -> dict:
+        ...
+
+    @abstractmethod
+    def session_param(self) -> dict:
+        ...
+
+    @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 PgDiskANNImplConfig(PgDiskANNIndexConfig):
+    index: IndexType = IndexType.DISKANN
+    max_neighbors: int | None
+    l_value_ib: int | None
+    l_value_is: float | None
+    maintenance_work_mem: Optional[str] = None
+    max_parallel_workers: Optional[int] = None
+
+    def index_param(self) -> dict:
+        return {
+            "metric": self.parse_metric(),
+            "index_type": self.index.value,
+            "options": {
+                "max_neighbors": self.max_neighbors,
+                "l_value_ib": self.l_value_ib,
+            },
+            "maintenance_work_mem": self.maintenance_work_mem,
+            "max_parallel_workers": self.max_parallel_workers,
+        }
+    
+    def search_param(self) -> dict:
+        return {
+            "metric": self.parse_metric(),
+            "metric_fun_op": self.parse_metric_fun_op(),
+        }
+    
+    def session_param(self) -> dict:
+        return {
+            "diskann.l_value_is": self.l_value_is,
+        }
+    
+_pgdiskann_case_config = {
+    IndexType.DISKANN: PgDiskANNImplConfig,
+}