From b72399ecad00eb7c02a846da3603da1f380c0c96 Mon Sep 17 00:00:00 2001 From: Murali Krishna <9718411+muralikpbhat@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:44:14 +0530 Subject: [PATCH 1/2] feat: add Infino client --- pyproject.toml | 1 + tests/test_infino.py | 68 +++++ vectordb_bench/backend/clients/__init__.py | 16 ++ vectordb_bench/backend/clients/infino/cli.py | 106 ++++++++ .../backend/clients/infino/config.py | 81 ++++++ .../backend/clients/infino/infino.py | 253 ++++++++++++++++++ vectordb_bench/cli/vectordbbench.py | 2 + vectordb_bench/frontend/config/styles.py | 1 + 8 files changed, 528 insertions(+) create mode 100644 tests/test_infino.py create mode 100644 vectordb_bench/backend/clients/infino/cli.py create mode 100644 vectordb_bench/backend/clients/infino/config.py create mode 100644 vectordb_bench/backend/clients/infino/infino.py diff --git a/pyproject.toml b/pyproject.toml index 223291721..b56eec93a 100644 --- a/pyproject.toml +++ b/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" ] diff --git a/tests/test_infino.py b/tests/test_infino.py new file mode 100644 index 000000000..3a58bf6b5 --- /dev/null +++ b/tests/test_infino.py @@ -0,0 +1,68 @@ +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(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}" diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index beacb37af..57f0ee429 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/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 diff --git a/vectordb_bench/backend/clients/infino/cli.py b/vectordb_bench/backend/clients/infino/cli.py new file mode 100644 index 000000000..3d181a90d --- /dev/null +++ b/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, + ) diff --git a/vectordb_bench/backend/clients/infino/config.py b/vectordb_bench/backend/clients/infino/config.py new file mode 100644 index 000000000..98a3d7e46 --- /dev/null +++ b/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 {} diff --git a/vectordb_bench/backend/clients/infino/infino.py b/vectordb_bench/backend/clients/infino/infino.py new file mode 100644 index 000000000..630759a6b --- /dev/null +++ b/vectordb_bench/backend/clients/infino/infino.py @@ -0,0 +1,253 @@ +import logging +import os +import tempfile +from collections.abc import Iterable +from contextlib import contextmanager +from pathlib import Path + +import infino +import numpy as np +import pyarrow as pa + +from ..api import VectorDB +from .config import InfinoIndexConfig + +log = logging.getLogger(__name__) + +_VECTOR_FIELD = "emb" +_ID_FIELD = "id" + + +class Infino(VectorDB): + """VectorDBBench client for Infino, an embedded vector/search engine. + + Infino is in-process: each benchmark worker connects to the same on-disk + catalog. The instance holds only picklable config so it survives the + ProcessPoolExecutor(spawn) boundary; the connection and table are opened + lazily in init(). + """ + + # Serialize the load: concurrent writes to a single table are not supported. + thread_safe: bool = False + + def __init__( + self, + dim: int, + db_config: dict, + db_case_config: InfinoIndexConfig, + collection_name: str = "vdbbench_infino", + drop_old: bool = False, + **kwargs, + ): + self.name = "Infino" + self.dim = dim + self.data_path = db_config["data_path"] + # A cache budget without a cache dir is a silent no-op in the + # engine (no disk cache is created); default the cache next to the + # catalog so warm queries are actually warm. + if db_config.get("cache_budget_bytes") and not db_config.get("cache_dir"): + db_config = {**db_config, "cache_dir": str(Path(self.data_path) / "cache")} + # Connection tuning (cache budget, cache dir, object-store options); pass only what is set. + self._connect_opts = { + k: db_config[k] + for k in ("cache_budget_bytes", "cache_dir", "storage_options") + if db_config.get(k) is not None + } + self.table_name = collection_name + self.metric = db_case_config.index_param()["metric"] + # Vector serving path + serve-time beam, bridged to the engine config + # before connect (see _apply_search_mode_config). + self._search_mode = db_case_config.search_mode + self._ef = db_case_config.ef + + self._conn = None + self._table = None + # _id -> dataset-id arrays; None until built (init() skips the build + # when the table is empty — the load-phase open) and after unpickling. + self._map_keys = None + self._map_vals = None + # Build the schema once so table creation and every append stay in lockstep. + self._schema = self._build_schema() + + # Open the connection once and keep it: create/drop the table here, and + # reuse the same handle in init() (reopening is costly and can deadlock). + # __getstate__ drops it so a spawned worker reopens its own. + Path(self.data_path).mkdir(parents=True, exist_ok=True) + self._conn = self._connect() + if drop_old: + # A drop invalidates the persisted _id map: a fresh ingest reassigns + # engine _ids, so the map (trusted whenever its row count matches) + # would be wrongly reused after re-ingesting the same number of rows. + self._id_map_path().unlink(missing_ok=True) + if self.table_name in self._conn.list_tables(): + self._conn.drop_table(self.table_name, purge=True) + if self.table_name not in self._conn.list_tables(): + self._conn.create_table(self.table_name, self._schema, self._index_spec()) + + def _apply_search_mode_config(self): + """Bridge ``search_mode`` and the serve-time beam to the engine's config. + + The engine reads ``vector.search_mode`` and ``vector.hnsw_ef_search`` from + ``$XDG_CONFIG_HOME/infino/config.yaml``; ``connect()`` has no equivalent + keywords, so a config file is the only way to select them without touching + the engine's public API. The engine loads that config lazily and caches + it for the process's lifetime, so the ``XDG_CONFIG_HOME`` override must + persist (it cannot be restored right after connect) — acceptable here + because each benchmark worker is a dedicated infino process. The engine + default is ``ivf`` with the stamped k->ef curve, so only values that + diverge from that are written and a pure-default run is byte-for-byte + unchanged. Idempotent; re-applied in each spawned worker before its first + connect. + """ + mode = self._search_mode + ef = self._ef or 0 + write_mode = bool(mode) and mode != "ivf" + if not write_mode and ef <= 0: + return + lines = ["vector:"] + if write_mode: + lines.append(f" search_mode: {mode}") + if ef > 0: + lines.append(f" hnsw_ef_search: {ef}") + cfg_root = Path(self.data_path) / "_infino_engine_cfg" + (cfg_root / "infino").mkdir(parents=True, exist_ok=True) + (cfg_root / "infino" / "config.yaml").write_text("\n".join(lines) + "\n") + os.environ["XDG_CONFIG_HOME"] = str(cfg_root) + + def _connect(self): + self._apply_search_mode_config() + return infino.connect(self.data_path, **self._connect_opts) + + def __getstate__(self) -> dict: + # Drop the non-picklable live connection so the instance can cross a process boundary. + return {**self.__dict__, "_conn": None, "_table": None, "_map_keys": None, "_map_vals": None} + + def _build_schema(self) -> pa.Schema: + return pa.schema( + [ + pa.field(_ID_FIELD, pa.int64(), nullable=False), + pa.field(_VECTOR_FIELD, pa.list_(pa.float32(), self.dim), nullable=False), + ] + ) + + def _index_spec(self) -> infino.IndexSpec: + return infino.IndexSpec().vector(_VECTOR_FIELD, self.dim, self.metric) + + @contextmanager + def init(self): + # Reuse one connection for the whole process: reopening is costly and can + # deadlock. __init__ opens it in the constructing process; a spawned + # worker (unpickled with _conn=None) opens its own here, once. + if self._conn is None: + self._conn = self._connect() + if self._table is None: + self._table = self._conn.open_table(self.table_name) + self._load_or_build_id_map() + yield + + # _id -> dataset-id translation, the same build-once / persist / reload + # pattern the engine's own bench uses for its ground-truth bin: one scan + # per TABLE (not per process), stored beside the catalog, reloaded by + # every worker in well under a second. _ids are 128-bit decimals, so the + # sorted key array is 16-byte big-endian bytes (lexicographic == numeric). + # + # The build lives in init() ON PURPOSE: init() runs before the timed + # search loops, so the scan/reload never lands inside a measured query. + # Two guards keep that placement safe: an empty table (the loader opens + # the connection before inserting a single row) builds and persists + # NOTHING, and a cached map is only trusted if its row count matches the + # table — otherwise it is rebuilt in place. + def _id_map_path(self) -> Path: + return Path(self.data_path) / f"{self.table_name}.idmap.npz" + + def _table_row_count(self) -> int: + res = self._conn.query_sql(f"SELECT COUNT(*) FROM {self.table_name}") + return int(res.column(0).to_pylist()[0]) + + def _load_or_build_id_map(self) -> None: + n_rows = self._table_row_count() + if n_rows == 0: + # Load-phase init(): drop_old has just recreated the table and no + # rows exist yet. Persisting an empty map here would poison every + # later search process (the cache is trusted once written). + return + path = self._id_map_path() + if path.exists(): + data = np.load(path) + if len(data["keys"]) == n_rows: + self._map_keys, self._map_vals = data["keys"], data["vals"] + return + # Row count mismatch: the cache belongs to a previous incarnation + # of the table (dropped and reloaded at a different size). Fall + # through and rebuild; os.replace keeps concurrent rebuilds safe. + m = self._conn.query_sql(f"SELECT _id, {_ID_FIELD} FROM {self.table_name}") + keys = np.array( + [int(v).to_bytes(16, "big") for v in m.column("_id").to_pylist()], + dtype="S16", + ) + vals = np.array(m.column(_ID_FIELD).to_pylist(), dtype=np.int64) + order = np.argsort(keys) + keys, vals = keys[order], vals[order] + # Suffix must end in .npz or np.savez appends it and orphans the file. + fd, tmp = tempfile.mkstemp(dir=path.parent, suffix=".tmp.npz") + os.close(fd) + np.savez(tmp, keys=keys, vals=vals) + Path(tmp).replace(path) + self._map_keys, self._map_vals = keys, vals + + def _to_dataset_ids(self, stable_ids: list) -> list[int]: + if not stable_ids: + return [] + if self._map_keys is None: + # Only reachable in a process that opened the connection while the + # table was still empty (load) and then searched (read-write + # cases). Ordinary search processes built the map in init(). + self._load_or_build_id_map() + if self._map_keys is None: + msg = ( + f"table {self.table_name!r} has no rows to map; the load stage " + "did not run (or wrote nothing) before search" + ) + raise RuntimeError(msg) + q = np.array([int(v).to_bytes(16, "big") for v in stable_ids], dtype="S16") + n = len(self._map_keys) + idx = np.searchsorted(self._map_keys, q) + # Every returned _id must be a key we mapped; a miss means the cached + # map belongs to a different table state — fail loudly, wrong ids + # here silently corrupt recall. + if (idx >= n).any() or (self._map_keys[idx.clip(max=n - 1)] != q).any(): + msg = ( + f"search returned _ids absent from the id map for {self.table_name!r}; " + f"stale cache at {self._id_map_path()} — delete it and rerun" + ) + raise RuntimeError(msg) + return self._map_vals[idx].tolist() + + def insert_embeddings( + self, + embeddings: Iterable[list[float]], + metadata: list[int], + **kwargs, + ) -> tuple[int, Exception | None]: + try: + arrays = [ + pa.array(metadata, type=pa.int64()), + pa.array(embeddings, type=pa.list_(pa.float32(), self.dim)), + ] + self._table.append(pa.record_batch(arrays, schema=self._schema)) + except Exception as e: + log.exception("Failed to insert embeddings into Infino") + return 0, e + return len(metadata), None + + def search_embedding(self, query: list[float], k: int = 100, **kwargs) -> list[int]: + # Vector serving is engine-decided; the call carries no tuning kwargs. + hits = self._table.vector_search(_VECTOR_FIELD, query, k) + return self._to_dataset_ids(hits.column("_id").to_pylist()) + + def optimize(self, data_size: int | None = None): + with self.init(): + self._table.optimize() + + def need_normalize_cosine(self) -> bool: + return True diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index e0cb98652..d42ab8d9f 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -14,6 +14,7 @@ ) from ..backend.clients.endee.cli import Endee from ..backend.clients.hologres.cli import HologresHGraph +from ..backend.clients.infino.cli import Infino from ..backend.clients.lancedb.cli import ( LanceDB, LanceDBAutoIndex, @@ -101,6 +102,7 @@ cli.add_command(Chroma) cli.add_command(Zvec) cli.add_command(Endee) +cli.add_command(Infino) cli.add_command(LindormIVFPQ) cli.add_command(LindormHNSW) cli.add_command(LindormIVFBQ) diff --git a/vectordb_bench/frontend/config/styles.py b/vectordb_bench/frontend/config/styles.py index 4b162e9ed..0592ca55d 100644 --- a/vectordb_bench/frontend/config/styles.py +++ b/vectordb_bench/frontend/config/styles.py @@ -81,6 +81,7 @@ def getPatternShape(i): # RedisCloud color: #0D6EFD # Chroma color: #FFC107 COLOR_MAP = { + DB.Infino.value: "#E8384F", DB.Milvus.value: "#0DCAF0", DB.ZillizCloud.value: "#0D6EFD", DB.ElasticCloud.value: "#04D6C8", From f18a433a44a3ef9bdb68cbc305a9bb012806bf40 Mon Sep 17 00:00:00 2001 From: Murali Krishna <9718411+muralikpbhat@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:22:01 +0530 Subject: [PATCH 2/2] infino: batch inserts into large appends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit VectorDBBench feeds inserts in small batches (default 100 rows) and each append() commits a superfile, so a large load fragmented into thousands of tiny superfiles — slow to load and slow to optimize. Buffer fed rows and commit them as one combined append at 100k rows, flushing any remainder when the load's init() scope exits (the same subprocess that inserted them), so a corpus smaller than the threshold is still fully persisted. Insert speed is now independent of the harness batch size; the served index is unchanged. Also fix the ef case-config test, which called index_param() without a metric_type and raised before reaching its assertion. --- tests/test_infino.py | 42 ++++++++++++- .../backend/clients/infino/infino.py | 61 +++++++++++++++++-- 2 files changed, 98 insertions(+), 5 deletions(-) diff --git a/tests/test_infino.py b/tests/test_infino.py index 3a58bf6b5..c468318a9 100644 --- a/tests/test_infino.py +++ b/tests/test_infino.py @@ -26,7 +26,7 @@ def test_ef_config(self): # 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(ef=768).index_param() + assert "ef" not in InfinoIndexConfig(metric_type=MetricType.COSINE, ef=768).index_param() with pytest.raises(ValueError, match="ef"): InfinoIndexConfig(ef=-1) @@ -66,3 +66,43 @@ def test_insert_and_search(self): 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]}" diff --git a/vectordb_bench/backend/clients/infino/infino.py b/vectordb_bench/backend/clients/infino/infino.py index 630759a6b..def210563 100644 --- a/vectordb_bench/backend/clients/infino/infino.py +++ b/vectordb_bench/backend/clients/infino/infino.py @@ -17,6 +17,16 @@ _VECTOR_FIELD = "emb" _ID_FIELD = "id" +# VectorDBBench feeds rows in small batches (its default is 100 rows), and each +# Infino append() commits a superfile. Committing one superfile per fed batch +# would fragment a large load into thousands of tiny superfiles, making both the +# load and the following optimize pathologically slow. insert_embeddings instead +# buffers fed rows and commits them as one combined append once this many have +# accumulated; the remainder is flushed when the load's init() scope exits, so a +# corpus smaller than this threshold is still fully persisted. The result is a +# handful of large superfiles regardless of the harness batch size. +_FLUSH_ROWS = 100_000 + class Infino(VectorDB): """VectorDBBench client for Infino, an embedded vector/search engine. @@ -66,6 +76,10 @@ def __init__( # when the table is empty — the load-phase open) and after unpickling. self._map_keys = None self._map_vals = None + # Rows accumulated by insert_embeddings as Arrow batches, committed as one + # large append at _FLUSH_ROWS and when the load's init() scope exits. + self._buf_batches: list[pa.RecordBatch] = [] + self._buf_rows = 0 # Build the schema once so table creation and every append stay in lockstep. self._schema = self._build_schema() @@ -119,8 +133,19 @@ def _connect(self): return infino.connect(self.data_path, **self._connect_opts) def __getstate__(self) -> dict: - # Drop the non-picklable live connection so the instance can cross a process boundary. - return {**self.__dict__, "_conn": None, "_table": None, "_map_keys": None, "_map_vals": None} + # Drop the non-picklable live connection so the instance can cross a + # process boundary. The buffer is always empty at a process boundary + # (the load subprocess flushes at init() exit before returning), so it + # is reset rather than shipped. + return { + **self.__dict__, + "_conn": None, + "_table": None, + "_map_keys": None, + "_map_vals": None, + "_buf_batches": [], + "_buf_rows": 0, + } def _build_schema(self) -> pa.Schema: return pa.schema( @@ -143,7 +168,14 @@ def init(self): if self._table is None: self._table = self._conn.open_table(self.table_name) self._load_or_build_id_map() - yield + try: + yield + finally: + # Commit any rows still buffered from the load. This runs in the same + # (sub)process that inserted them, before it returns — the only point + # at which a load smaller than _FLUSH_ROWS would otherwise never be + # persisted. A no-op outside the load path (the buffer is empty). + self._flush() # _id -> dataset-id translation, the same build-once / persist / reload # pattern the engine's own bench uses for its ground-truth bin: one scan @@ -229,17 +261,38 @@ def insert_embeddings( metadata: list[int], **kwargs, ) -> tuple[int, Exception | None]: + # Buffer fed rows and commit them as a few large superfiles rather than + # one per call (see _FLUSH_ROWS); the remainder is flushed at init() exit + # so every fed row is persisted before search. try: arrays = [ pa.array(metadata, type=pa.int64()), pa.array(embeddings, type=pa.list_(pa.float32(), self.dim)), ] - self._table.append(pa.record_batch(arrays, schema=self._schema)) + self._buf_batches.append(pa.record_batch(arrays, schema=self._schema)) + self._buf_rows += len(metadata) + if self._buf_rows >= _FLUSH_ROWS: + self._flush() except Exception as e: log.exception("Failed to insert embeddings into Infino") return 0, e return len(metadata), None + def _flush(self) -> None: + """Commit all buffered rows as a single superfile and clear the buffer. + + Concatenates the buffered batches into one contiguous batch so the load + commits large superfiles. Inserts are serialized (the client is not + thread-safe, so the runner drives a single worker), so no lock is needed. + """ + if not self._buf_batches: + return + table = pa.Table.from_batches(self._buf_batches, schema=self._schema).combine_chunks() + for batch in table.to_batches(): + self._table.append(batch) + self._buf_batches = [] + self._buf_rows = 0 + def search_embedding(self, query: list[float], k: int = 100, **kwargs) -> list[int]: # Vector serving is engine-decided; the call carries no tuning kwargs. hits = self._table.vector_search(_VECTOR_FIELD, query, k)