Add a VectorDBBench client for Infino (vector search)

pyproject.toml

@@ -82,6 +82,7 @@ doris           = [ "doris-vector-search" ]
 turbopuffer     = [ "turbopuffer" ]
 zvec            = [ "zvec" ]
 endee           = [ "endee==0.1.10" ]
+infino          = [ "infino>=0.5.6" ]
 lindorm         = [ "opensearch-py" ]
 seekdb          = [ "mysql-connector-python" ]
 volc_mysql      = [ "mysql-connector-python" ]

tests/test_infino.py

@@ -0,0 +1,108 @@
+import tempfile
+
+import numpy as np
+import pytest
+
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.backend.clients.api import MetricType
+from vectordb_bench.backend.clients.infino.config import InfinoIndexConfig
+
+
+class TestInfino:
+    def test_search_mode_config(self):
+        # Default is the resident-HNSW path; ivf is opt-out; bad values reject.
+        # search_mode is bridged to the engine config file, so it must NOT leak
+        # into index_param() (which feeds IndexSpec).
+        assert InfinoIndexConfig().search_mode == "hnsw_ivf"
+        cfg = InfinoIndexConfig(metric_type=MetricType.COSINE, search_mode="ivf")
+        assert cfg.search_mode == "ivf"
+        assert "search_mode" not in cfg.index_param()
+        with pytest.raises(ValueError, match="search_mode"):
+            InfinoIndexConfig(search_mode="bogus")
+
+    def test_ef_config(self):
+        # ef defaults to 0 (serve at the graph's stamped k->ef curve); a positive
+        # value is a fixed serve-time beam; negative is rejected. Like search_mode
+        # it is bridged to the engine config, not index_param().
+        assert InfinoIndexConfig().ef == 0
+        assert InfinoIndexConfig(ef=768).ef == 768
+        assert "ef" not in InfinoIndexConfig(metric_type=MetricType.COSINE, ef=768).index_param()
+        with pytest.raises(ValueError, match="ef"):
+            InfinoIndexConfig(ef=-1)
+
+    def test_insert_and_search(self):
+        pytest.importorskip("infino")  # engine binding; skip if the extra isn't installed
+        assert DB.Infino.value == "Infino"
+
+        dbcls = DB.Infino.init_cls
+        config_cls = DB.Infino.config_cls
+        case_config_cls = DB.Infino.case_config_cls()
+
+        dim = 16
+        count = 2_000
+        rng = np.random.default_rng(0)
+        embeddings = rng.random((count, dim)).tolist()
+
+        with tempfile.TemporaryDirectory() as data_path:
+            db_config = config_cls(data_path=data_path).to_dict()
+            # 2K rows sit inside the engine's default rerank budget, so the
+            # engine-decided serving is exact for the assertion.
+            db_case_config = case_config_cls(metric_type=MetricType.L2)
+
+            client = dbcls(
+                dim=dim,
+                db_config=db_config,
+                db_case_config=db_case_config,
+                collection_name="test_infino",
+                drop_old=True,
+            )
+
+            with client.init():
+                inserted, err = client.insert_embeddings(embeddings=embeddings, metadata=list(range(count)))
+                assert err is None
+                assert inserted == count
+
+            with client.init():
+                test_id = 42
+                res = client.search_embedding(query=embeddings[test_id], k=10)
+                assert res[0] == test_id, f"nearest neighbor id {res[0]} != query id {test_id}"
+
+    def test_insert_buffering_persists_every_row(self, monkeypatch: pytest.MonkeyPatch):
+        # insert_embeddings buffers fed rows and commits large appends. Rows must
+        # survive both the mid-load threshold flush and the init()-exit flush of
+        # the sub-threshold remainder (the case where the corpus < _FLUSH_ROWS).
+        pytest.importorskip("infino")  # engine binding; skip if the extra isn't installed
+        monkeypatch.setattr("vectordb_bench.backend.clients.infino.infino._FLUSH_ROWS", 50)
+
+        dim = 16  # engine requires dim in [16, 4096]
+        count = 130  # 20-row feeds flush at 60 twice (120 rows), leaving a 10-row remainder
+        rng = np.random.default_rng(1)
+        embeddings = rng.random((count, dim)).tolist()
+
+        dbcls = DB.Infino.init_cls
+        config_cls = DB.Infino.config_cls
+        case_config_cls = DB.Infino.case_config_cls()
+
+        with tempfile.TemporaryDirectory() as data_path:
+            client = dbcls(
+                dim=dim,
+                db_config=config_cls(data_path=data_path).to_dict(),
+                db_case_config=case_config_cls(metric_type=MetricType.L2),
+                collection_name="test_infino_buffer",
+                drop_old=True,
+            )
+            with client.init():
+                total = 0
+                for start in range(0, count, 20):
+                    chunk = embeddings[start : start + 20]
+                    inserted, err = client.insert_embeddings(
+                        embeddings=chunk, metadata=list(range(start, start + len(chunk)))
+                    )
+                    assert err is None
+                    total += inserted
+                assert total == count
+            # The 10-row remainder was flushed at init() exit; a row from it
+            # (id 129, the last inserted) must be searchable.
+            with client.init():
+                res = client.search_embedding(query=embeddings[129], k=1)
+                assert res[0] == 129, f"remainder row not persisted: got {res[0]}"

vectordb_bench/backend/clients/__init__.py

@@ -65,6 +65,7 @@ class DB(Enum):
     SeekDB = "SeekDB"
     VolcMySQL = "VolcMySQL"
     Adbpg = "AnalyticDB for PostgreSQL"
+    Infino = "Infino"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
@@ -281,6 +282,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return Adbpg
 
+        if self == DB.Infino:
+            from .infino.infino import Infino
+
+            return Infino
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -499,6 +505,11 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return AdbpgConfig
 
+        if self == DB.Infino:
+            from .infino.config import InfinoConfig
+
+            return InfinoConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -719,6 +730,11 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912, PLR0915
 
             return AdbpgIndexConfig
 
+        if self == DB.Infino:
+            from .infino.config import InfinoIndexConfig
+
+            return InfinoIndexConfig
+
         # DB.Pinecone, DB.Redis
         return EmptyDBCaseConfig
 

vectordb_bench/backend/clients/infino/cli.py

@@ -0,0 +1,106 @@
+from typing import Annotated, Unpack
+
+import click
+
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+
+from .config import _DEFAULT_CACHE_BUDGET_BYTES
+
+DBTYPE = DB.Infino
+
+
+def _parse_kv_list(_ctx, _param, values) -> dict[str, str]:  # noqa: ANN001
+    """Parse repeatable or comma-separated key=value items into a dict."""
+    parsed: dict[str, str] = {}
+    for item in values or ():
+        for part in (p.strip() for p in str(item).split(",")):
+            if not part:
+                continue
+            if "=" not in part:
+                msg = f"Expect key=value, got: {part}"
+                raise click.BadParameter(msg)
+            k, v = part.split("=", 1)
+            parsed[k.strip()] = v.strip()
+    return parsed
+
+
+# Shared connection options for the vector command.
+_data_path_option = click.option(
+    "--data-path", type=str, default="/tmp/vectordb_bench/infino", help="Infino catalog directory"
+)
+_cache_budget_option = click.option(
+    "--cache-budget-bytes",
+    type=int,
+    default=_DEFAULT_CACHE_BUDGET_BYTES,
+    help="Disk-cache ceiling in bytes; raise for corpora larger than the cache",
+)
+_cache_dir_option = click.option("--cache-dir", type=str, default=None, help="Infino disk-cache directory")
+_storage_option_option = click.option(
+    "--storage-option",
+    "storage_options",
+    type=str,
+    multiple=True,
+    callback=_parse_kv_list,
+    help="Object-store option as key=value (repeatable or comma-separated), e.g. region=us-east-1",
+)
+
+
+class InfinoCommonTypedDict(CommonTypedDict):
+    data_path: Annotated[str, _data_path_option]
+    cache_budget_bytes: Annotated[int, _cache_budget_option]
+    cache_dir: Annotated[str, _cache_dir_option]
+    storage_options: Annotated[dict, _storage_option_option]
+
+
+class InfinoTypedDict(InfinoCommonTypedDict):
+    table_name: Annotated[
+        str,
+        click.option("--table-name", type=str, default="vdbbench_infino", help="Infino table name"),
+    ]
+    search_mode: Annotated[
+        str,
+        click.option(
+            "--search-mode",
+            type=click.Choice(["ivf", "hnsw_ivf"]),
+            default="hnsw_ivf",
+            help="Vector serving path, bridged to the engine config: "
+            "hnsw_ivf (default; resident HNSW graph with ivf fallback) or ivf",
+        ),
+    ]
+    ef: Annotated[
+        int,
+        click.option(
+            "--ef",
+            type=int,
+            default=0,
+            help="Serve-time HNSW beam (search_mode=hnsw_ivf), bridged to "
+            "vector.hnsw_ef_search. 0 (default) uses the graph's stamped k->ef "
+            "curve; a positive value fixes the beam. Sweep it across runs (one "
+            "value per run) to trace the recall/latency curve without a rebuild.",
+        ),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(InfinoTypedDict)
+def Infino(**parameters: Unpack[InfinoTypedDict]):
+    from .config import InfinoConfig, InfinoIndexConfig
+
+    run(
+        db=DBTYPE,
+        db_config=InfinoConfig(
+            data_path=parameters["data_path"],
+            table_name=parameters["table_name"],
+            cache_budget_bytes=parameters["cache_budget_bytes"],
+            cache_dir=parameters["cache_dir"],
+            storage_options=parameters["storage_options"] or None,
+        ),
+        db_case_config=InfinoIndexConfig(search_mode=parameters["search_mode"], ef=parameters["ef"]),
+        **parameters,
+    )

vectordb_bench/backend/clients/infino/config.py

@@ -0,0 +1,81 @@
+from pydantic import BaseModel, field_validator
+
+from vectordb_bench.backend.clients.api import DBCaseConfig, DBConfig, MetricType
+
+# Vector serving path, bridged to the engine's `vector.search_mode` config key
+# by the client before connect (see infino.py) — a config override, not a public
+# engine-API field. "hnsw_ivf" (default) builds + serves a resident HNSW graph
+# over the quantized vectors (automatic IVF fallback) — the path this client
+# benchmarks; "ivf" serves the reclaimable IVF scan.
+_SEARCH_MODES = ("ivf", "hnsw_ivf")
+
+# Infino distance metrics; all are distances where smaller means nearer.
+_METRIC_MAP = {
+    MetricType.COSINE: "cosine",
+    MetricType.L2: "l2sq",
+    MetricType.IP: "negdot",
+}
+
+
+# Disk-cache ceiling, not a preallocation: sized well above the 10 GiB engine
+# default so large corpora stay cached instead of falling back to range-only reads.
+_DEFAULT_CACHE_BUDGET_BYTES = 64 * 1024**3
+
+
+class InfinoConfig(DBConfig):
+    data_path: str = "/tmp/vectordb_bench/infino"
+    table_name: str = "vdbbench_infino"
+    cache_budget_bytes: int = _DEFAULT_CACHE_BUDGET_BYTES
+    cache_dir: str | None = None
+    storage_options: dict[str, str] | None = None
+
+    def to_dict(self) -> dict:
+        return {
+            "data_path": self.data_path,
+            "table_name": self.table_name,
+            "cache_budget_bytes": self.cache_budget_bytes,
+            "cache_dir": self.cache_dir,
+            "storage_options": self.storage_options,
+        }
+
+
+class InfinoIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    # Serving path + beam are forwarded through the engine config file, not
+    # IndexSpec (see _SEARCH_MODES above). Default hnsw_ivf: the resident-HNSW
+    # path is what this client benchmarks.
+    search_mode: str = "hnsw_ivf"
+    # Serve-time HNSW beam for search_mode=hnsw_ivf, bridged to the engine's
+    # vector.hnsw_ef_search config key (see infino.py). 0 (default) serves each
+    # query at the graph's stamped k->ef curve; a positive value fixes the beam,
+    # so sweeping ef across runs traces the recall/QPS curve of one built graph
+    # with no rebuild.
+    ef: int = 0
+
+    @field_validator("search_mode")
+    @classmethod
+    def _validate_search_mode(cls, v: str) -> str:
+        if v not in _SEARCH_MODES:
+            msg = f"Infino search_mode must be one of {_SEARCH_MODES}, got {v!r}"
+            raise ValueError(msg)
+        return v
+
+    @field_validator("ef")
+    @classmethod
+    def _validate_ef(cls, v: int) -> int:
+        if v < 0:
+            msg = f"Infino ef must be >= 0 (0 = use the stamped curve), got {v}"
+            raise ValueError(msg)
+        return v
+
+    def parse_metric(self) -> str:
+        if self.metric_type not in _METRIC_MAP:
+            msg = f"Infino does not support metric {self.metric_type}"
+            raise ValueError(msg)
+        return _METRIC_MAP[self.metric_type]
+
+    def index_param(self) -> dict:
+        return {"metric": self.parse_metric()}
+
+    def search_param(self) -> dict:
+        return {}