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..c468318a9 --- /dev/null +++ b/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]}" 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..def210563 --- /dev/null +++ b/vectordb_bench/backend/clients/infino/infino.py @@ -0,0 +1,306 @@ +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" + +# 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. + + 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 + # 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() + + # 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. 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( + [ + 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() + 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 + # 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]: + # 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._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) + 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",