vectordb-bench 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

vectordb_bench/backend/clients/oss_opensearch/cli.py

@@ -0,0 +1,155 @@
+import logging
+from typing import Annotated, TypedDict, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor1,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+from .config import OSSOpenSearchQuantization, OSSOS_Engine
+
+log = logging.getLogger(__name__)
+
+
+class OSSOpenSearchTypedDict(TypedDict):
+    host: Annotated[str, click.option("--host", type=str, help="Db host", required=True)]
+    port: Annotated[int, click.option("--port", type=int, default=80, help="Db Port")]
+    user: Annotated[str, click.option("--user", type=str, help="Db User")]
+    password: Annotated[str, click.option("--password", type=str, help="Db password")]
+    number_of_shards: Annotated[
+        int,
+        click.option("--number-of-shards", type=int, help="Number of primary shards for the index", default=1),
+    ]
+    number_of_replicas: Annotated[
+        int,
+        click.option(
+            "--number-of-replicas", type=int, help="Number of replica copies for each primary shard", default=1
+        ),
+    ]
+    index_thread_qty: Annotated[
+        int,
+        click.option(
+            "--index-thread-qty",
+            type=int,
+            help="Thread count for native engine indexing",
+            default=4,
+        ),
+    ]
+
+    engine: Annotated[
+        str,
+        click.option(
+            "--engine",
+            type=click.Choice(["nmslib", "faiss", "lucene"], case_sensitive=False),
+            help="HNSW algorithm implementation to use",
+            default="faiss",
+        ),
+    ]
+
+    metric_type: Annotated[
+        str,
+        click.option(
+            "--metric-type",
+            type=click.Choice(["l2", "cosine", "ip"], case_sensitive=False),
+            help="Distance metric type for vector similarity",
+            default="l2",
+        ),
+    ]
+
+    number_of_segments: Annotated[
+        int,
+        click.option("--number-of-segments", type=int, help="Target number of segments after merging", default=1),
+    ]
+
+    refresh_interval: Annotated[
+        str,
+        click.option(
+            "--refresh-interval", type=str, help="How often to make new data available for search", default="60s"
+        ),
+    ]
+
+    force_merge_enabled: Annotated[
+        bool,
+        click.option("--force-merge-enabled", type=bool, help="Whether to perform force merge operation", default=True),
+    ]
+
+    flush_threshold_size: Annotated[
+        str,
+        click.option(
+            "--flush-threshold-size", type=str, help="Size threshold for flushing the transaction log", default="5120mb"
+        ),
+    ]
+
+    cb_threshold: Annotated[
+        str,
+        click.option(
+            "--cb-threshold",
+            type=str,
+            help="k-NN Memory circuit breaker threshold",
+            default="50%",
+        ),
+    ]
+
+    quantization_type: Annotated[
+        str | None,
+        click.option(
+            "--quantization-type",
+            type=click.Choice(["fp32", "fp16"]),
+            help="quantization type for vectors (in index)",
+            default="fp32",
+            required=False,
+        ),
+    ]
+
+    engine: Annotated[
+        str | None,
+        click.option(
+            "--engine",
+            type=click.Choice(["faiss", "lucene"]),
+            help="quantization type for vectors (in index)",
+            default="faiss",
+            required=False,
+        ),
+    ]
+
+
+class OSSOpenSearchHNSWTypedDict(CommonTypedDict, OSSOpenSearchTypedDict, HNSWFlavor1): ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(OSSOpenSearchHNSWTypedDict)
+def OSSOpenSearch(**parameters: Unpack[OSSOpenSearchHNSWTypedDict]):
+    from .config import OSSOpenSearchConfig, OSSOpenSearchIndexConfig
+
+    run(
+        db=DB.OSSOpenSearch,
+        db_config=OSSOpenSearchConfig(
+            host=parameters["host"],
+            port=parameters["port"],
+            user=parameters["user"],
+            password=SecretStr(parameters["password"]),
+        ),
+        db_case_config=OSSOpenSearchIndexConfig(
+            number_of_shards=parameters["number_of_shards"],
+            number_of_replicas=parameters["number_of_replicas"],
+            index_thread_qty=parameters["index_thread_qty"],
+            number_of_segments=parameters["number_of_segments"],
+            refresh_interval=parameters["refresh_interval"],
+            force_merge_enabled=parameters["force_merge_enabled"],
+            flush_threshold_size=parameters["flush_threshold_size"],
+            index_thread_qty_during_force_merge=parameters["index_thread_qty_during_force_merge"],
+            cb_threshold=parameters["cb_threshold"],
+            efConstruction=parameters["ef_construction"],
+            efSearch=parameters["ef_runtime"],
+            M=parameters["m"],
+            engine=OSSOS_Engine(parameters["engine"]),
+            quantization_type=OSSOpenSearchQuantization(parameters["quantization_type"]),
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/oss_opensearch/config.py

@@ -0,0 +1,157 @@
+import logging
+from enum import Enum
+
+from pydantic import BaseModel, SecretStr, root_validator, validator
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+log = logging.getLogger(__name__)
+
+
+class OSSOpenSearchConfig(DBConfig, BaseModel):
+    host: str = ""
+    port: int = 80
+    user: str | None = None
+    password: SecretStr | None = None
+
+    def to_dict(self) -> dict:
+        use_ssl = self.port == 443
+        http_auth = (
+            (self.user, self.password.get_secret_value())
+            if self.user is not None and self.password is not None and len(self.user) != 0 and len(self.password) != 0
+            else ()
+        )
+        return {
+            "hosts": [{"host": self.host, "port": self.port}],
+            "http_auth": http_auth,
+            "use_ssl": use_ssl,
+            "http_compress": True,
+            "verify_certs": use_ssl,
+            "ssl_assert_hostname": False,
+            "ssl_show_warn": False,
+            "timeout": 600,
+        }
+
+    @validator("*")
+    def not_empty_field(cls, v: any, field: any):
+        if (
+            field.name in cls.common_short_configs()
+            or field.name in cls.common_long_configs()
+            or field.name in ["user", "password", "host"]
+        ):
+            return v
+        if isinstance(v, str | SecretStr) and len(v) == 0:
+            raise ValueError("Empty string!")
+        return v
+
+
+class OSSOS_Engine(Enum):
+    faiss = "faiss"
+    lucene = "lucene"
+
+
+class OSSOpenSearchQuantization(Enum):
+    fp32 = "fp32"
+    fp16 = "fp16"
+
+
+class OSSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType = MetricType.L2
+    engine: OSSOS_Engine = OSSOS_Engine.faiss
+    efConstruction: int = 256
+    efSearch: int = 100
+    engine_name: str | None = None
+    metric_type_name: str | None = None
+    M: int = 16
+    index_thread_qty: int | None = 4
+    number_of_shards: int | None = 1
+    number_of_replicas: int | None = 0
+    number_of_segments: int | None = 1
+    refresh_interval: str | None = "60s"
+    force_merge_enabled: bool | None = True
+    flush_threshold_size: str | None = "5120mb"
+    index_thread_qty_during_force_merge: int = 8
+    cb_threshold: str | None = "50%"
+    number_of_indexing_clients: int | None = 1
+    use_routing: bool = False  # for label-filter cases
+    oversample_factor: float = 1.0
+    quantization_type: OSSOpenSearchQuantization = OSSOpenSearchQuantization.fp32
+
+    @root_validator
+    def validate_engine_name(cls, values: dict):
+        """Map engine_name string from UI to engine enum"""
+        if values.get("engine_name"):
+            engine_name = values["engine_name"].lower()
+            if engine_name == "faiss":
+                values["engine"] = OSSOS_Engine.faiss
+            elif engine_name == "lucene":
+                values["engine"] = OSSOS_Engine.lucene
+            else:
+                log.warning(f"Unknown engine_name: {engine_name}, defaulting to faiss")
+                values["engine"] = OSSOS_Engine.faiss
+        return values
+
+    def __eq__(self, obj: any):
+        return (
+            self.engine == obj.engine
+            and self.M == obj.M
+            and self.efConstruction == obj.efConstruction
+            and self.number_of_shards == obj.number_of_shards
+            and self.number_of_replicas == obj.number_of_replicas
+            and self.number_of_segments == obj.number_of_segments
+            and self.use_routing == obj.use_routing
+            and self.quantization_type == obj.quantization_type
+        )
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                self.engine,
+                self.M,
+                self.efConstruction,
+                self.number_of_shards,
+                self.number_of_replicas,
+                self.number_of_segments,
+                self.use_routing,
+                self.quantization_type,
+            )
+        )
+
+    def parse_metric(self) -> str:
+        log.info(f"User specified metric_type: {self.metric_type_name}")
+        self.metric_type = MetricType[self.metric_type_name.upper()]
+        if self.metric_type == MetricType.IP:
+            return "innerproduct"
+        if self.metric_type == MetricType.COSINE:
+            return "cosinesimil"
+        if self.metric_type == MetricType.L2:
+            log.info("Using l2 as specified by user")
+            return "l2"
+        return "l2"
+
+    @property
+    def use_quant(self) -> bool:
+        return self.quantization_type is not OSSOpenSearchQuantization.fp32
+
+    def index_param(self) -> dict:
+        log.info(f"Using engine: {self.engine} for index creation")
+        log.info(f"Using metric_type: {self.metric_type_name} for index creation")
+        log.info(f"Resulting space_type: {self.parse_metric()} for index creation")
+
+        return {
+            "name": "hnsw",
+            "engine": self.engine.value,
+            "space_type": self.parse_metric(),
+            "parameters": {
+                "ef_construction": self.efConstruction,
+                "m": self.M,
+                **(
+                    {"encoder": {"name": "sq", "parameters": {"type": self.quantization_type.value}}}
+                    if self.use_quant
+                    else {}
+                ),
+            },
+        }
+
+    def search_param(self) -> dict:
+        return {"ef_search": self.efSearch}