Add vespa integration

Author: nuvotex-tk

README.md

@@ -56,6 +56,7 @@ All the database client supported
 | aliyun_opensearch        | `pip install vectordb-bench[aliyun_opensearch]` |
 | mongodb                  | `pip install vectordb-bench[mongodb]`       |
 | tidb                     | `pip install vectordb-bench[tidb]`          |
+| vespa                    | `pip install vectordb-bench[vespa]`         |
 
 ### Run
 

install/requirements_py3.11.txt

@@ -23,3 +23,4 @@ pydantic<v2
 scikit-learn
 pymilvus
 clickhouse_connect
+pyvespa

pyproject.toml

@@ -70,7 +70,8 @@ all = [
     "alibabacloud_searchengine20211025",
     "mariadb",
     "PyMySQL",
-    "clickhouse-connect"
+    "clickhouse-connect",
+    "pyvespa",
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -92,6 +93,7 @@ mongodb         = [ "pymongo" ]
 mariadb         = [ "mariadb" ]
 tidb            = [ "PyMySQL" ]
 clickhouse      = [ "clickhouse-connect" ]
+vespa           = [ "pyvespa" ]
 
 [project.urls]
 "repository" = "https://github.com/zilliztech/VectorDBBench"

vectordb_bench/backend/clients/__init__.py

@@ -44,6 +44,7 @@ class DB(Enum):
     MongoDB = "MongoDB"
     TiDB = "TiDB"
     Clickhouse = "Clickhouse"
+    Vespa = "Vespa"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
@@ -157,6 +158,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
             from .test.test import Test
 
             return Test
+        
+        if self == DB.Vespa:
+            from .vespa.vespa import Vespa
+
+            return Vespa
 
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
@@ -273,6 +279,12 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901
             from .test.config import TestConfig
 
             return TestConfig
+        
+        if self == DB.Vespa:
+            from .vespa.config import VespaConfig
+
+            return VespaConfig
+
 
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
@@ -365,6 +377,11 @@ def case_config_cls(  # noqa: PLR0911
             from .tidb.config import TiDBIndexConfig
 
             return TiDBIndexConfig
+        
+        if self == DB.Vespa:
+            from .vespa.config import VespaHNSWConfig
+
+            return VespaHNSWConfig
 
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig

vectordb_bench/backend/clients/vespa/cli.py

@@ -0,0 +1,47 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    HNSWFlavor1,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+
+
+class VespaTypedDict(CommonTypedDict, HNSWFlavor1):
+    uri: Annotated[
+        str,
+        click.option("--uri", "-u", type=str, help="uri connection string", default="http://127.0.0.1"),
+    ]
+    port: Annotated[
+        int,
+        click.option("--port", "-p", type=int, help="connection port", default=8080),
+    ]
+    quantization: Annotated[
+        str, click.option("--quantization", type=click.Choice(["none", "binary"], case_sensitive=False), default="none")
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(VespaTypedDict)
+def Vespa(**params: Unpack[VespaTypedDict]):
+    from .config import VespaConfig, VespaHNSWConfig
+
+    case_params = {
+        "quantization_type": params["quantization"],
+        "M": params["m"],
+        "efConstruction": params["ef_construction"],
+        "ef": params["ef_search"],
+    }
+
+    run(
+        db=DB.Vespa,
+        db_config=VespaConfig(url=SecretStr(params["uri"]), port=params["port"]),
+        db_case_config=VespaHNSWConfig(**{k: v for k, v in case_params.items() if v}),
+        **params,
+    )

vectordb_bench/backend/clients/vespa/config.py

@@ -0,0 +1,51 @@
+from typing import Literal, TypeAlias
+
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, MetricType
+
+VespaMetric: TypeAlias = Literal["euclidean", "angular", "dotproduct", "prenormalized-angular", "hamming", "geodegrees"]
+
+VespaQuantizationType: TypeAlias = Literal["none", "binary"]
+
+
+class VespaConfig(DBConfig):
+    url: SecretStr = "http://127.0.0.1"
+    port: int = 8080
+
+    def to_dict(self):
+        return {
+            "url": self.url.get_secret_value(),
+            "port": self.port,
+        }
+
+
+class VespaHNSWConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType = MetricType.COSINE
+    quantization_type: VespaQuantizationType = "none"
+    M: int = 16
+    efConstruction: int = 200
+    ef: int = 100
+
+    def index_param(self) -> dict:
+        return {
+            "distance_metric": self.parse_metric(self.metric_type),
+            "max_links_per_node": self.M,
+            "neighbors_to_explore_at_insert": self.efConstruction,
+        }
+
+    def search_param(self) -> dict:
+        return {}
+
+    def parse_metric(self, metric_type: MetricType) -> VespaMetric:
+        match metric_type:
+            case MetricType.COSINE:
+                return "angular"
+            case MetricType.L2:
+                return "euclidean"
+            case MetricType.DP | MetricType.IP:
+                return "dotproduct"
+            case MetricType.HAMMING:
+                return "hamming"
+            case _:
+                raise NotImplementedError

vectordb_bench/backend/clients/vespa/util.py

@@ -0,0 +1,16 @@
+"""Utility functions for supporting binary quantization
+
+From https://docs.vespa.ai/en/binarizing-vectors.html#appendix-conversion-to-int8
+"""
+import numpy as np
+
+
+def binarize_tensor(tensor: list[float]) -> list[int]:
+    """
+    Binarize a floating-point list by thresholding at zero
+    and packing the bits into bytes.
+    """
+    tensor = np.array(tensor)
+    return (
+        np.packbits(np.where(tensor > 0, 1, 0), axis=0).astype(np.int8).tolist()
+    )