diff --git a/Makefile b/Makefile index 5ec0917ec..e227e3b92 100644 --- a/Makefile +++ b/Makefile @@ -1,89 +1,8 @@ -PYTHON ?= /home/rowan/.local/share/uv/tools/pip/bin/python -BENCH = $(PYTHON) -m vectordb_bench.cli.vectordbbench -COMMON = --drop-old --skip-search-concurrent -ANTFLY_NUM_PER_BATCH ?= 5000 - -help: - @echo "VectorDBBench — Antfly Integration" - @echo "" - @echo "Benchmarks (50K vectors, 1536 dim):" - @echo " make bench-antfly-50k Antfly SPANN" - @echo " make bench-qdrant-50k Qdrant HNSW" - @echo " make bench-milvus-50k Milvus AutoIndex" - @echo " make bench-all-50k All three (50K)" - @echo "" - @echo "Benchmarks (1M vectors, 768 dim):" - @echo " make bench-antfly-1m Antfly SPANN" - @echo " make bench-qdrant-1m Qdrant HNSW" - @echo " make bench-milvus-1m Milvus AutoIndex" - @echo " make bench-all-1m All three (1M)" - @echo "" - @echo "Infrastructure:" - @echo " make start-qdrant Start Qdrant (Docker)" - @echo " make start-milvus Start Milvus (Docker)" - @echo " make stop-all Stop Docker containers" - @echo "" - @echo "Dev:" - @echo " make unittest Run unit tests" - @echo " make format / lint Code formatting" - -# --- 50K benchmarks --- -bench-antfly-50k: - NUM_PER_BATCH=$(ANTFLY_NUM_PER_BATCH) $(BENCH) antflyaknn --host localhost --port 8080 --num-shards 1 \ - --case-type Performance1536D50K $(COMMON) --db-label antfly-local - -bench-qdrant-50k: - $(BENCH) qdrantlocal --url http://localhost:6333 \ - --case-type Performance1536D50K $(COMMON) --db-label qdrant-local \ - --m 16 --ef-construct 128 --hnsw-ef 64 - -bench-milvus-50k: - $(BENCH) milvusautoindex --uri http://localhost:19530 \ - --case-type Performance1536D50K $(COMMON) --db-label milvus-local - -bench-all-50k: bench-antfly-50k bench-qdrant-50k bench-milvus-50k - -# --- 1M benchmarks --- -bench-antfly-1m: - NUM_PER_BATCH=$(ANTFLY_NUM_PER_BATCH) $(BENCH) antflyaknn --host localhost --port 8080 --num-shards 1 \ - --case-type Performance768D1M $(COMMON) --db-label antfly-local - -bench-qdrant-1m: - $(BENCH) qdrantlocal --url http://localhost:6333 \ - --case-type Performance768D1M $(COMMON) --db-label qdrant-local \ - --m 16 --ef-construct 128 --hnsw-ef 64 - -bench-milvus-1m: - $(BENCH) milvusautoindex --uri http://localhost:19530 \ - --case-type Performance768D1M $(COMMON) --db-label milvus-local - -bench-all-1m: bench-antfly-1m bench-qdrant-1m bench-milvus-1m - -# --- Infrastructure --- -start-qdrant: - docker run -d --name qdrant_bench -p 6333:6333 -p 6334:6334 qdrant/qdrant:latest - -start-milvus: - docker run -d --name milvus_bench -p 19530:19530 -p 9091:9091 \ - -e ETCD_USE_EMBED=true -e COMMON_STORAGETYPE=local \ - milvusdb/milvus:latest milvus run standalone - -stop-all: - -docker stop qdrant_bench milvus_bench 2>/dev/null - -docker rm qdrant_bench milvus_bench 2>/dev/null - -# --- Dev --- -unittest: - PYTHONPATH=`pwd` python3 -m pytest tests/test_dataset.py::TestDataSet::test_download_small -svv - -format: - PYTHONPATH=`pwd` python3 -m black vectordb_bench - PYTHONPATH=`pwd` python3 -m ruff check vectordb_bench --fix - lint: PYTHONPATH=`pwd` python3 -m black vectordb_bench --check PYTHONPATH=`pwd` python3 -m ruff check vectordb_bench -.PHONY: help bench-antfly-50k bench-qdrant-50k bench-milvus-50k bench-all-50k \ - bench-antfly-1m bench-qdrant-1m bench-milvus-1m bench-all-1m \ - start-qdrant start-milvus stop-all unittest format lint +unittest: + PYTHONPATH=`pwd` python3 -m pytest tests/test_dataset.py::TestDataSet::test_download_small -svv + +.PHONY: lint unittest diff --git a/antfly-recall-antfly-only.md b/antfly-recall-antfly-only.md deleted file mode 100644 index 4c249f91b..000000000 --- a/antfly-recall-antfly-only.md +++ /dev/null @@ -1,176 +0,0 @@ -# Antfly Recall Assessment: Antfly-Only Findings - -## Scope - -This note is intentionally limited to **Antfly-side** explanations for the low recall seen in `#antfly > vector-db-bench`. - -I am **not** treating the original `VectorDBBench` client wiring issues as part of this diagnosis. Those confounders were removed before the rerun summarized below. The question here is narrower: after controlling for the obvious external issues, does the remaining low recall look like an Antfly problem? - -## Bottom line - -Yes, at this point the low recall looks much more like an **Antfly-side issue** than a benchmark wiring mistake. - -After rerunning with the benchmark-side confounders removed, recall stayed essentially unchanged: - -- `recall@100 = 0.5008` -- `ndcg = 0.6166` - -That means the remaining plausible causes are now mostly inside Antfly itself: - -1. a concrete cosine/HBC correctness bug visible in the source-built upstream server -2. hardcoded HBC search defaults that may be too aggressive for this workload -3. dynamic pruning behavior in HBC -4. quantized leaf search -5. possible drift between the running local Antfly binary and the checked-out source/API expectations - -## Findings - -### 1. The remaining problem is no longer well explained by client wiring - -I reran the 50K case after removing the main external confounders and the result remained about the same as before. That materially weakens the hypothesis that recall was low simply because the benchmark client was wired incorrectly. - -So the updated diagnosis is: - -- the original setup did have avoidable confounders -- but those confounders were **not** the main reason recall ended up near 50% - -That pushes the investigation onto Antfly itself. - -### 2. The source-built upstream Antfly still fails after receiving unit-normalized cosine vectors - -I then reran against a fresh source-built upstream checkout at `~/Documents/antfly/antfly-upstream` after hardening the client further: - -- correct cosine `distance_metric` -- explicit cosine normalization on insert/query in the Antfly client -- `need_normalize_cosine()` enabled so `VectorDBBench` normalizes cosine datasets before calling the client - -For the first insert batch, the benchmark logged: - -- `metric=COSINE` -- `use_cosine=True` -- `first_sqnorm=1.000000` - -So at that point the benchmark was definitely sending a unit-normalized cosine vector. - -Even with that in place, the upstream source build still panicked inside Antfly during HBC split/reclustering with: - -- `vector is not a unit vector` -- stack rooted in `BalancedBinaryKmeans.validateVectors()` via `HBCIndex.splitVectorSet()` - -That is the strongest current evidence in this whole investigation. It means there is at least one remaining Antfly-side cosine/HBC correctness problem that is **not** explained by the benchmark simply forgetting to normalize vectors. - -### 3. Antfly's dense-vector path is using hardcoded HBC defaults that could plausibly cap recall - -The embeddings index builds an HBC index with a fixed configuration in source: - -- quantization enabled -- `Episilon2 = 7` -- `BranchingFactor = 168` -- `LeafSize = 168` -- `SearchWidth = 1008` - -Relevant source: - -- [`src/store/db/indexes/embeddings_index.go:538`](/home/rowan/Documents/antfly/antfly/src/store/db/indexes/embeddings_index.go#L538) - -Those are not obviously conservative "maximize recall" settings. They look like fixed performance-oriented defaults. If recall is poor on this dataset, this config is a very credible place to look first. - -### 4. HBC is doing aggressive dynamic pruning during search - -The HBC search code prunes internal nodes and leaves based on the current best distance threshold and `Episilon2`: - -- internal-node pruning: - [`lib/vectorindex/hbc.go:1885`](/home/rowan/Documents/antfly/antfly/lib/vectorindex/hbc.go#L1885) -- default distance metric fallback: - [`lib/vectorindex/hbc.go:288`](/home/rowan/Documents/antfly/antfly/lib/vectorindex/hbc.go#L288) - -The specific pruning logic here is aggressive enough that it could absolutely explain a recall ceiling if the heuristic is not well tuned for this workload. - -I would treat this as one of the strongest Antfly-side suspects. - -### 5. Antfly is searching quantized vectors by default - -The embeddings index turns quantization on by default in the HBC config: - -- [`src/store/db/indexes/embeddings_index.go:544`](/home/rowan/Documents/antfly/antfly/src/store/db/indexes/embeddings_index.go#L544) - -And leaf search uses quantized vectors when available: - -- [`lib/vectorindex/hbc.go:1902`](/home/rowan/Documents/antfly/antfly/lib/vectorindex/hbc.go#L1902) - -That does not prove quantization is the reason recall is low, but it is another real Antfly-side approximation step in the retrieval path. Between quantization and dynamic pruning, there is enough approximation here that a ~50% recall result is not inherently surprising. - -### 6. The running local Antfly binary appears inconsistent with the checked-out source/API expectations - -This matters because it raises the possibility that the benchmark is not actually exercising the exact implementation implied by the local source tree. - -Observed on the running local Antfly used for the rerun: - -- creating the vector index with type `"embeddings"` failed with HTTP 400 -- creating it with `"aknn_v0"` succeeded -- inserts using `SyncLevelEmbeddings`/`aknn` semantics for precomputed `_embeddings` failed with a server-side `index not found: vec` -- explicitly including `"name": "vec"` in the index config was necessary for the field-only precomputed-embedding path to work reliably - -Those runtime behaviors are notable because the checked-out source clearly presents an embeddings index abstraction: - -- the API type includes `EmbeddingsIndexConfig` and documents `distance_metric`: - [`src/store/db/indexes/openapi.gen.go:170`](/home/rowan/Documents/antfly/antfly/src/store/db/indexes/openapi.gen.go#L170) -- the CLI expects an explicit `name` field in an index definition: - [`cmd/antfly/cmd/cli/table.go:133`](/home/rowan/Documents/antfly/antfly/cmd/antfly/cmd/cli/table.go#L133) -- the DB tests explicitly expect vector writes to support embeddings-level sync before search: - [`src/store/db/embeddings_test.go:792`](/home/rowan/Documents/antfly/antfly/src/store/db/embeddings_test.go#L792) -- `SyncLevelEmbeddings` is a first-class sync level: - [`src/store/db/ops.proto:41`](/home/rowan/Documents/antfly/antfly/src/store/db/ops.proto#L41) - -So there is at least some evidence of **runtime/build drift or compatibility mismatch**: - -- either the running binary is older or behaviorally different from the checked-out source -- or the field-only precomputed-embedding path is not as solid as the source/tests suggest - -Either way, that is an Antfly-side problem, not a benchmark problem. - -### 7. The direct query shape itself does not look like the main failure point - -The Antfly query path for direct embeddings appears coherent in source: - -- request embeddings are copied into the remote index query: - [`src/metadata/api_query.go:45`](/home/rowan/Documents/antfly/antfly/src/metadata/api_query.go#L45) -- if `indexes` is omitted, Antfly skips query-time embedding generation: - [`src/metadata/api_query.go:496`](/home/rowan/Documents/antfly/antfly/src/metadata/api_query.go#L496) -- remote search accepts direct embedding searches: - [`src/store/db/indexes/remoteindex.go:1430`](/home/rowan/Documents/antfly/antfly/src/store/db/indexes/remoteindex.go#L1430) - -So my current read is not "the query API is fundamentally broken." The stronger Antfly-side story is that the query reaches the right subsystem, but the subsystem's current implementation and/or runtime configuration is not producing good enough nearest-neighbor quality. - -## Most likely Antfly-side explanations, in order - -1. There is a cosine/HBC bug in Antfly's indexing path, visible even on the source-built upstream server after the benchmark sends unit-normalized vectors. -2. HBC pruning and search-width defaults are too aggressive for this benchmark. -3. Quantization is costing too much recall on this workload. -4. The running local Antfly binary does not match the behavior implied by the checked-out source, especially around embeddings index creation and sync semantics. -5. There is a field-only precomputed-embedding bug or edge case in the current implementation. - -## What I would test next in Antfly - -1. Reproduce the source-build cosine panic in a focused Antfly test around `HBCIndex.splitVectorSet()` / `BalancedBinaryKmeans.validateVectors()` with precomputed unit vectors. -2. Verify why Antfly later observes a non-unit vector after the benchmark has already sent `sqnorm=1.0`. -3. Verify the exact binary/version/config that was running for the local ~50% recall rerun. -4. Run the same 50K case with a less approximate Antfly configuration: - - higher `SearchWidth` - - lower or disabled pruning - - quantization disabled -5. Check whether the runtime binary is actually the same code as `~/Documents/antfly/antfly`. -6. Reproduce the `SyncLevelEmbeddings` failure in a focused Antfly test using precomputed `_embeddings`. - -## Verdict - -The general concern in the thread now looks directionally correct: the remaining low recall appears to be an **Antfly problem**. - -The specific root cause is not yet proven to be only "HBC/SPANN pruning," but the evidence now points much more strongly to: - -- an Antfly cosine/HBC bug on the source-built path -- Antfly search heuristics/defaults -- Antfly approximation behavior -- and possibly Antfly runtime/build inconsistency - -not to simple benchmark miswiring. diff --git a/pyproject.toml b/pyproject.toml index f19759c7a..c240ae55a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ "polars", "plotly", "environs", - "pydantic=2.0", "scikit-learn", "pymilvus", # with pandas, numpy, ujson "ujson", @@ -79,7 +79,7 @@ all = [ "mysql-connector-python", "turbopuffer[fast]", 'zvec', - "endee==0.1.10", # compatible with pydantic<2 + "endee==0.1.10", # TODO: check pydantic v2 compat ] qdrant = [ "qdrant-client" ] diff --git a/tests/test_bench_runner.py b/tests/test_bench_runner.py index 5fab91067..e65c73887 100644 --- a/tests/test_bench_runner.py +++ b/tests/test_bench_runner.py @@ -52,7 +52,7 @@ def test_performance_case_no_error(self): ) t = task_config.copy() - d = t.json(exclude={'db_config': {'password', 'api_key'}}) + d = t.model_dump_json(exclude={'db_config': {'password', 'api_key'}}) log.info(f"{d}") import ujson diff --git a/vectordb_bench/backend/clients/alisql/config.py b/vectordb_bench/backend/clients/alisql/config.py index f942c30f1..16c56e90f 100644 --- a/vectordb_bench/backend/clients/alisql/config.py +++ b/vectordb_bench/backend/clients/alisql/config.py @@ -49,8 +49,8 @@ def parse_metric(self) -> str: class AliSQLHNSWConfig(AliSQLIndexConfig, DBCaseConfig): - M: int | None - ef_search: int | None + M: int | None = None + ef_search: int | None = None index: IndexType = IndexType.HNSW def index_param(self) -> dict: diff --git a/vectordb_bench/backend/clients/alloydb/config.py b/vectordb_bench/backend/clients/alloydb/config.py index d6e54e487..11e65084e 100644 --- a/vectordb_bench/backend/clients/alloydb/config.py +++ b/vectordb_bench/backend/clients/alloydb/config.py @@ -43,8 +43,8 @@ class AlloyDBIndexParam(TypedDict): metric: str index_type: str index_creation_with_options: Sequence[dict[str, Any]] - maintenance_work_mem: str | None - max_parallel_workers: int | None + maintenance_work_mem: str | None = None + max_parallel_workers: int | None = None class AlloyDBSearchParam(TypedDict): @@ -120,15 +120,15 @@ def _optionally_build_set_options( class AlloyDBScaNNConfig(AlloyDBIndexConfig): index: IndexType = IndexType.SCANN - num_leaves: int | None - quantizer: str | None - enable_pca: str | None - max_num_levels: int | None - num_leaves_to_search: int | None - max_top_neighbors_buffer_size: int | None - pre_reordering_num_neighbors: int | None - num_search_threads: int | None - max_num_prefetch_datasets: int | None + num_leaves: int | None = None + quantizer: str | None = None + enable_pca: str | None = None + max_num_levels: int | None = None + num_leaves_to_search: int | None = None + max_top_neighbors_buffer_size: int | None = None + pre_reordering_num_neighbors: int | None = None + num_search_threads: int | None = None + max_num_prefetch_datasets: int | None = None maintenance_work_mem: str | None = None max_parallel_workers: int | None = None diff --git a/vectordb_bench/backend/clients/antfly/antfly.py b/vectordb_bench/backend/clients/antfly/antfly.py index a31a3a346..cd43b3edc 100644 --- a/vectordb_bench/backend/clients/antfly/antfly.py +++ b/vectordb_bench/backend/clients/antfly/antfly.py @@ -1,5 +1,8 @@ +import base64 import logging import math +import os +import struct import time from contextlib import contextmanager from typing import Any @@ -13,7 +16,7 @@ BATCH_CHUNK_SIZE = 500 TABLE_READY_TIMEOUT = 30 TABLE_READY_POLL_INTERVAL = 2 -INDEX_READY_TIMEOUT = 1800 +INDEX_READY_TIMEOUT = 7200 INDEX_READY_POLL_INTERVAL = 2 INDEX_NAME = "vec" INDEX_TYPES = ("embeddings", "aknn_v0") @@ -41,50 +44,68 @@ def __init__( self.collection_name = collection_name self.dim = dim - base_url = f"http://{_httpx_host(db_config['host'])}:{db_config['port']}/api/v1" - self._base_url = base_url + self._metadata_base_url = f"http://{_httpx_host(db_config['host'])}:{db_config['port']}/api/v1" + self._store_host = _httpx_host(db_config.get("store_host") or db_config["host"]) + self._store_port = db_config.get("store_port") + self._use_direct_store_search = bool(db_config.get("use_direct_store_search")) + self._pack_query_vectors = bool(db_config.get("pack_query_vectors")) + self._direct_shard_id: str | None = None num_shards = db_config.get("num_shards", 1) - client = httpx.Client(base_url=base_url, timeout=60) + if self._use_direct_store_search and not self._store_port: + raise ValueError("Antfly direct store search requires store_port to be configured") + + client = httpx.Client(base_url=self._metadata_base_url, timeout=60) try: if drop_old: r = client.delete(f"/tables/{self.collection_name}") log.info(f"Drop table response: {r.status_code}") - # Two-step table creation to avoid Pebble lock issue: - # 1. Create table without embeddings index - r = client.post(f"/tables/{self.collection_name}", json={"num_shards": num_shards}) - log.info(f"Create table response: {r.status_code}") + table = self._get_table_status_or_none(client) + if table is None: + r = client.post(f"/tables/{self.collection_name}", json={"num_shards": num_shards}) + log.info(f"Create table response: {r.status_code}") + r.raise_for_status() + else: + log.info("Reusing existing table: %s", self.collection_name) - # Wait for shard to initialize self._wait_for_shard_ready(client) - # 2. Add field-only embeddings index (pre-computed vectors, no embedder needed) - index_def = { - "name": INDEX_NAME, - "dimension": dim, - "external": True, - **self.case_config.index_param(), - } - index_error = None - for index_type in INDEX_TYPES: - r = client.post( - f"/tables/{self.collection_name}/indexes/{INDEX_NAME}", - json={"type": index_type, **index_def}, - ) - log.info(f"Add embeddings index response ({index_type}): {r.status_code}") - if r.is_success: - index_error = None - break - index_error = r - if index_error is not None: - index_error.raise_for_status() + if self._get_index_status(client) is None: + index_def = { + "name": INDEX_NAME, + "dimension": dim, + "external": True, + **self.case_config.index_param(), + } + index_error = None + # Try each index type, with and without field, to handle + # both old binaries (require field) and new source (reject field with external). + for index_type in INDEX_TYPES: + for extra in ({}, {"field": SOURCE_FIELD}): + r = client.post( + f"/tables/{self.collection_name}/indexes/{INDEX_NAME}", + json={"type": index_type, **index_def, **extra}, + ) + log.info( + f"Add embeddings index response ({index_type}, field={'field' in extra}): {r.status_code}" + ) + if r.is_success: + index_error = None + break + index_error = r + if index_error is None: + break + if index_error is not None: + index_error.raise_for_status() + else: + log.info("Reusing existing embeddings index: %s", INDEX_NAME) self._wait_for_index_ready(client, expected_total=0) + self._refresh_direct_search_routing(client) finally: client.close() def _wait_for_shard_ready(self, client: httpx.Client): - """Wait for shard to be initialized and accepting writes.""" deadline = time.monotonic() + TABLE_READY_TIMEOUT while time.monotonic() < deadline: try: @@ -93,7 +114,6 @@ def _wait_for_shard_ready(self, client: httpx.Client): json={"inserts": {"_healthcheck": {"_probe": True}}, "sync_level": "write"}, ) if r.status_code < 500: - # Delete the probe doc client.post( f"/tables/{self.collection_name}/batch", json={"deletes": ["_healthcheck"], "sync_level": "write"}, @@ -112,6 +132,28 @@ def _get_index_status(self, client: httpx.Client) -> dict | None: r.raise_for_status() return r.json() + def _get_table_status(self, client: httpx.Client) -> dict: + r = client.get(f"/tables/{self.collection_name}") + r.raise_for_status() + return r.json() + + def _get_table_status_or_none(self, client: httpx.Client) -> dict | None: + r = client.get(f"/tables/{self.collection_name}") + if r.status_code == 404: + return None + r.raise_for_status() + return r.json() + + def _refresh_direct_search_routing(self, client: httpx.Client): + if not self._use_direct_store_search: + return + table = self._get_table_status(client) + shards = table.get("shards") or {} + if len(shards) != 1: + msg = f"Antfly direct store search currently requires exactly one shard; found {len(shards)} shards" + raise ValueError(msg) + self._direct_shard_id = next(iter(shards)) + def _index_status_is_ready( self, payload: dict | None, @@ -157,17 +199,26 @@ def _wait_for_index_ready(self, client: httpx.Client, expected_total: int | None @contextmanager def init(self): - self.client = httpx.Client(base_url=self._base_url, timeout=120) + self.client = httpx.Client(base_url=self._metadata_base_url, timeout=120) + self.store_client = None try: + if self._use_direct_store_search: + self.store_client = httpx.Client(base_url=self._store_base_url, timeout=120) yield finally: self.client.close() self.client = None + if self.store_client is not None: + self.store_client.close() + self.store_client = None + + @property + def _store_base_url(self) -> str: + if self._store_port is None: + raise ValueError("Antfly store_base_url requested without store_port configured") + return f"http://{self._store_host}:{self._store_port}" def need_normalize_cosine(self) -> bool: - # TaskRunner already gates normalization on the dataset metric being COSINE. - # Returning True here avoids relying on mutable case_config state to decide - # whether Antfly should receive unit-normalized vectors. return True def _uses_cosine_distance(self) -> bool: @@ -183,11 +234,73 @@ def _normalize_vector(vector: list[float]) -> list[float]: return vector return [value / norm for value in vector] + @staticmethod + def _pack_vector(vector: list[float]) -> str: + raw = struct.pack(f"<{len(vector)}f", *vector) + return base64.b64encode(raw).decode("ascii") + + def _serialize_query_vector(self, vector: list[float]) -> list[float] | str: + if self._pack_query_vectors or os.environ.get("ANTFLY_PACK_VECTORS") == "1": + return self._pack_vector(vector) + return vector + + def _metadata_query_body(self, query: list[float], k: int) -> dict[str, Any]: + return { + "embeddings": {"vec": self._serialize_query_vector(query)}, + "limit": k, + "fields": ["id"], + **self.case_config.search_param(), + } + + def _store_query_body(self, query: list[float], k: int) -> dict[str, Any]: + search_params = self.case_config.search_param() + vector_paging_options: dict[str, Any] = {"limit": k} + if "search_effort" in search_params: + vector_paging_options["search_effort"] = search_params["search_effort"] + return { + "star": True, + "limit": k, + "vector_searches": {INDEX_NAME: self._serialize_query_vector(query)}, + "vector_paging_options": vector_paging_options, + } + + def _parse_metadata_hits(self, data: dict) -> list[int]: + resp = data.get("responses", [{}])[0] + hits_obj = resp.get("hits") or {} + hits = hits_obj.get("hits") or [] + results = [] + for hit in hits: + if "id" in hit: + results.append(int(hit["id"])) + else: + doc_key = hit.get("_id", "") + try: + results.append(int(doc_key.split(":", 1)[1])) + except (IndexError, ValueError): + log.warning(f"Could not parse id from _id: {doc_key}") + return results + + def _parse_store_hits(self, data: dict) -> list[int]: + vec_result = (data.get("search_result") or {}).get(INDEX_NAME) or {} + hits = vec_result.get("hits") or [] + results = [] + for hit in hits: + fields = hit.get("fields") or {} + if "id" in fields: + results.append(int(fields["id"])) + continue + doc_key = hit.get("id", "") + try: + results.append(int(doc_key.split(":", 1)[1])) + except (IndexError, ValueError): + log.warning(f"Could not parse id from direct-store hit id: {doc_key}") + return results + def ready_to_search(self) -> bool: if getattr(self, "client", None) is not None: payload = self._get_index_status(self.client) return self._index_status_is_ready(payload, payload.get("status") if payload else None) - with httpx.Client(base_url=self._base_url, timeout=120) as client: + with httpx.Client(base_url=self._metadata_base_url, timeout=120) as client: payload = self._get_index_status(client) return self._index_status_is_ready(payload, payload.get("status") if payload else None) @@ -195,7 +308,7 @@ def optimize(self, data_size: int | None = None): if getattr(self, "client", None) is not None: self._wait_for_index_ready(self.client, expected_total=data_size) return - with httpx.Client(base_url=self._base_url, timeout=120) as client: + with httpx.Client(base_url=self._metadata_base_url, timeout=120) as client: self._wait_for_index_ready(client, expected_total=data_size) def insert_embeddings( @@ -215,27 +328,15 @@ def insert_embeddings( embedding = embeddings[i] if use_cosine: embedding = self._normalize_vector(embedding) + serialized_embedding = self._serialize_query_vector(embedding) inserts[key] = { "id": metadata[i], "metadata": metadata[i], - # Antfly derives a hashID for precomputed embeddings from the - # configured source field. Give it a stable string value instead - # of forcing the "missing field" path for every document. SOURCE_FIELD: str(metadata[i]), - "_embeddings": {"vec": embedding}, + "_embeddings": {"vec": serialized_embedding}, } - payload = {"inserts": inserts, "sync_level": "aknn"} + payload = {"inserts": inserts, "sync_level": "write"} r = self.client.post(f"/tables/{self.collection_name}/batch", json=payload) - if not r.is_success: - log.warning( - "Antfly aknn batch write failed (%s), falling back to sync_level=write: %s", - r.status_code, - r.text, - ) - r = self.client.post( - f"/tables/{self.collection_name}/batch", - json={"inserts": inserts, "sync_level": "write"}, - ) r.raise_for_status() except Exception as e: log.warning(f"Antfly insert error: {e}") @@ -250,33 +351,23 @@ def search_embedding( timeout: int | None = None, **kwargs: Any, ) -> list[int]: - # Workaround: omit semantic_search and indexes, send embeddings directly. - # This bypasses generateQueryEmbeddings() and uses the direct vector path. if self._uses_cosine_distance(): query = self._normalize_vector(query) - body = { - "embeddings": {"vec": query}, - "limit": k, - "fields": ["id"], - } + + if self._use_direct_store_search: + if self._direct_shard_id is None: + self._refresh_direct_search_routing(self.client) + r = self.store_client.post( + "/search", + headers={"X-Raft-Shard-Id": self._direct_shard_id}, + json=self._store_query_body(query, k), + ) + r.raise_for_status() + return self._parse_store_hits(r.json()) + r = self.client.post( f"/tables/{self.collection_name}/query", - json=body, + json=self._metadata_query_body(query, k), ) r.raise_for_status() - data = r.json() - # Response format: {"responses": [{"hits": {"hits": [{"_id": "key:42", ...}, ...]}}]} - resp = data.get("responses", [{}])[0] - hits_obj = resp.get("hits") or {} - hits = hits_obj.get("hits") or [] - results = [] - for hit in hits: - if "id" in hit: - results.append(int(hit["id"])) - else: - doc_key = hit.get("_id", "") - try: - results.append(int(doc_key.split(":", 1)[1])) - except (IndexError, ValueError): - log.warning(f"Could not parse id from _id: {doc_key}") - return results + return self._parse_metadata_hits(r.json()) diff --git a/vectordb_bench/backend/clients/antfly/cli.py b/vectordb_bench/backend/clients/antfly/cli.py index 4410d638b..4587dae0f 100644 --- a/vectordb_bench/backend/clients/antfly/cli.py +++ b/vectordb_bench/backend/clients/antfly/cli.py @@ -13,13 +13,45 @@ class AntflyTypedDict(TypedDict): - host: Annotated[str, click.option("--host", type=str, help="Antfly host", default="localhost", show_default=True)] - port: Annotated[int, click.option("--port", type=int, help="Antfly port", default=8080, show_default=True)] + host: Annotated[ + str, click.option("--host", type=str, help="Antfly metadata API host", default="localhost", show_default=True) + ] + port: Annotated[ + int, click.option("--port", type=int, help="Antfly metadata API port", default=8080, show_default=True) + ] + store_host: Annotated[str, click.option("--store-host", type=str, help="Antfly store API host", default=None)] + store_port: Annotated[ + int, click.option("--store-port", type=int, help="Antfly store API port for direct search", default=None) + ] username: Annotated[str, click.option("--username", type=str, help="Antfly username", default=None)] password: Annotated[str, click.option("--password", type=str, help="Antfly password", default=None)] num_shards: Annotated[ int, click.option("--num-shards", type=int, help="Number of shards", default=1, show_default=True) ] + use_direct_store_search: Annotated[ + bool, + click.option( + "--use-direct-store-search/--no-direct-store-search", + help="Query Antfly through the store /search API instead of the metadata table query API", + default=False, + show_default=True, + ), + ] + pack_query_vectors: Annotated[ + bool, + click.option( + "--pack-query-vectors/--no-pack-query-vectors", + help="Send query vectors in Antfly's packed base64 float32 wire format", + default=False, + show_default=True, + ), + ] + search_effort: Annotated[ + float, + click.option( + "--search-effort", type=float, help="Search effort 0.0-1.0 (higher=better recall, slower)", default=None + ), + ] class AntflyAKNNTypedDict(CommonTypedDict, AntflyTypedDict): ... @@ -36,12 +68,17 @@ def AntflyAKNN(**parameters: Unpack[AntflyAKNNTypedDict]): db_label=parameters["db_label"], host=parameters["host"], port=parameters["port"], + store_host=parameters["store_host"], + store_port=parameters["store_port"], username=SecretStr(parameters["username"]) if parameters["username"] else None, password=SecretStr(parameters["password"]) if parameters["password"] else None, num_shards=parameters["num_shards"], + use_direct_store_search=parameters["use_direct_store_search"], + pack_query_vectors=parameters["pack_query_vectors"], ), db_case_config=AntflyIndexConfig( num_shards=parameters["num_shards"], + search_effort=parameters["search_effort"], ), **parameters, ) diff --git a/vectordb_bench/backend/clients/antfly/config.py b/vectordb_bench/backend/clients/antfly/config.py index db3c6c809..f83c96062 100644 --- a/vectordb_bench/backend/clients/antfly/config.py +++ b/vectordb_bench/backend/clients/antfly/config.py @@ -6,23 +6,32 @@ class AntflyConfig(DBConfig): host: str = "localhost" port: int = 8080 + store_host: str | None = None + store_port: int | None = None username: SecretStr | None = None password: SecretStr | None = None num_shards: int = 1 + use_direct_store_search: bool = False + pack_query_vectors: bool = False def to_dict(self) -> dict: return { "host": self.host, "port": self.port, + "store_host": self.store_host, + "store_port": self.store_port, "username": self.username.get_secret_value() if self.username else None, "password": self.password.get_secret_value() if self.password else None, "num_shards": self.num_shards, + "use_direct_store_search": self.use_direct_store_search, + "pack_query_vectors": self.pack_query_vectors, } class AntflyIndexConfig(BaseModel, DBCaseConfig): metric_type: MetricType | None = None num_shards: int = 1 + search_effort: float | None = None def parse_metric(self) -> str: if self.metric_type == MetricType.COSINE: @@ -35,4 +44,6 @@ def index_param(self) -> dict: return {"distance_metric": self.parse_metric()} def search_param(self) -> dict: + if self.search_effort is not None: + return {"search_effort": self.search_effort} return {} diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index 82eda1824..abd9e1c22 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -1,8 +1,9 @@ from abc import ABC, abstractmethod from contextlib import contextmanager from enum import StrEnum +from typing import Any -from pydantic import BaseModel, SecretStr, validator +from pydantic import BaseModel, ConfigDict, model_validator from vectordb_bench.backend.filter import Filter, FilterOp @@ -67,6 +68,8 @@ class DBConfig(ABC, BaseModel): ZillizCloudConfig.db_label = 1cu-perf """ + model_config = ConfigDict(validate_default=True) + db_label: str = "" version: str = "" note: str = "" @@ -89,13 +92,18 @@ def common_long_configs() -> list[str]: def to_dict(self) -> dict: raise NotImplementedError - @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(): - return v - if not v and isinstance(v, str | SecretStr): - raise ValueError("Empty string!") - return v + @model_validator(mode="before") + @classmethod + def not_empty_field(cls, data: Any) -> Any: + if isinstance(data, dict): + skip = set(cls.common_short_configs()) | set(cls.common_long_configs()) + for name, v in data.items(): + if name in skip: + continue + if isinstance(v, str) and len(v) == 0: + msg = f"Empty string for field '{name}'!" + raise ValueError(msg) + return data class DBCaseConfig(ABC): diff --git a/vectordb_bench/backend/clients/aws_opensearch/config.py b/vectordb_bench/backend/clients/aws_opensearch/config.py index 5ab63010d..f6c81923d 100644 --- a/vectordb_bench/backend/clients/aws_opensearch/config.py +++ b/vectordb_bench/backend/clients/aws_opensearch/config.py @@ -1,7 +1,8 @@ import logging from enum import Enum +from typing import Any -from pydantic import BaseModel, SecretStr, validator +from pydantic import BaseModel, SecretStr, model_validator from ..api import DBCaseConfig, DBConfig, MetricType @@ -32,17 +33,18 @@ def to_dict(self) -> dict: "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 + @model_validator(mode="before") + @classmethod + def not_empty_field(cls, data: Any) -> Any: + if isinstance(data, dict): + skip = set(cls.common_short_configs()) | set(cls.common_long_configs()) | {"user", "password", "host"} + for name, v in data.items(): + if name in skip: + continue + if isinstance(v, str) and len(v) == 0: + msg = f"Empty string for field '{name}'!" + raise ValueError(msg) + return data class AWSOS_Engine(Enum): diff --git a/vectordb_bench/backend/clients/chroma/config.py b/vectordb_bench/backend/clients/chroma/config.py index cd3e01ecc..8c1cac286 100644 --- a/vectordb_bench/backend/clients/chroma/config.py +++ b/vectordb_bench/backend/clients/chroma/config.py @@ -6,7 +6,7 @@ class ChromaConfig(DBConfig): user: str | None = None - password: SecretStr | None + password: SecretStr | None = None host: SecretStr = "localhost" port: int = 8000 @@ -26,7 +26,7 @@ def to_dict(self) -> dict: class ChromaIndexConfig(ChromaConfig, DBCaseConfig): - metric_type: MetricType = "cosine" + metric_type: MetricType = MetricType.COSINE m: int = 16 ef_construct: int = 100 ef_search: int | None = 100 diff --git a/vectordb_bench/backend/clients/cockroachdb/config.py b/vectordb_bench/backend/clients/cockroachdb/config.py index 0d608da8f..88ec0e5ea 100644 --- a/vectordb_bench/backend/clients/cockroachdb/config.py +++ b/vectordb_bench/backend/clients/cockroachdb/config.py @@ -75,16 +75,16 @@ class CockroachDBIndexParam(TypedDict): metric: str index_creation_with_options: Sequence[dict[str, Any]] - min_partition_size: int | None - max_partition_size: int | None - build_beam_size: int | None + min_partition_size: int | None = None + max_partition_size: int | None = None + build_beam_size: int | None = None class CockroachDBSearchParam(TypedDict): """Search parameters for CockroachDB vector queries.""" metric_fun_op: LiteralString - vector_search_beam_size: int | None + vector_search_beam_size: int | None = None class CockroachDBSessionCommands(TypedDict): diff --git a/vectordb_bench/backend/clients/doris/config.py b/vectordb_bench/backend/clients/doris/config.py index a15309922..9caf72eb3 100644 --- a/vectordb_bench/backend/clients/doris/config.py +++ b/vectordb_bench/backend/clients/doris/config.py @@ -1,6 +1,7 @@ import logging +from typing import Any -from pydantic import BaseModel, SecretStr, validator +from pydantic import BaseModel, SecretStr, model_validator from ..api import DBCaseConfig, DBConfig, MetricType @@ -17,9 +18,10 @@ class DorisConfig(DBConfig): db_name: str = "test" ssl: bool = False - @validator("*") - def not_empty_field(cls, v: any, field: any): - return v + @model_validator(mode="before") + @classmethod + def not_empty_field(cls, data: Any) -> Any: + return data # allow all fields including empty strings def to_dict(self) -> dict: pwd_str = self.password.get_secret_value() diff --git a/vectordb_bench/backend/clients/lindorm/config.py b/vectordb_bench/backend/clients/lindorm/config.py index 367f369e3..0e0d4aae6 100644 --- a/vectordb_bench/backend/clients/lindorm/config.py +++ b/vectordb_bench/backend/clients/lindorm/config.py @@ -43,9 +43,9 @@ def parse_metric(self) -> str: class HNSWConfig(LindormIndexConfig, DBCaseConfig): index: IndexType = IndexType.HNSW - M: int | None - efConstruction: int | None - efSearch: int | None + M: int | None = None + efConstruction: int | None = None + efSearch: int | None = None filter_type: str | None = "efficient_filter" k_expand_scope: int | None = 1000 @@ -72,12 +72,12 @@ def search_param(self, do_filter: bool = False) -> dict: # first layer searching for cluster centroids is hnsw class IVFPQConfig(LindormIndexConfig, DBCaseConfig): index: IndexType = IndexType.IVFPQ - nlist: int | None - nprobe: int | None + nlist: int | None = None + nprobe: int | None = None # search parameters - centroids_hnsw_M: int | None - centroids_hnsw_efConstruction: int | None - centroids_hnsw_efSearch: int | None + centroids_hnsw_M: int | None = None + centroids_hnsw_efConstruction: int | None = None + centroids_hnsw_efSearch: int | None = None filter_type: str | None = "efficient_filter" reorder_factor: int | None = 10 @@ -116,13 +116,13 @@ def search_param(self, do_filter: bool = False) -> dict: class IVFBQConfig(LindormIndexConfig, DBCaseConfig): index: IndexType = IndexType.IVFBQ - nlist: int | None - exbits: int | None - nprobe: int | None + nlist: int | None = None + exbits: int | None = None + nprobe: int | None = None # search parameters - centroids_hnsw_M: int | None - centroids_hnsw_efConstruction: int | None - centroids_hnsw_efSearch: int | None + centroids_hnsw_M: int | None = None + centroids_hnsw_efConstruction: int | None = None + centroids_hnsw_efSearch: int | None = None filter_type: str | None = "efficient_filter" reorder_factor: int | None = 10 diff --git a/vectordb_bench/backend/clients/mariadb/config.py b/vectordb_bench/backend/clients/mariadb/config.py index d183adc76..21ea9ac2e 100644 --- a/vectordb_bench/backend/clients/mariadb/config.py +++ b/vectordb_bench/backend/clients/mariadb/config.py @@ -46,11 +46,11 @@ def parse_metric(self) -> str: class MariaDBHNSWConfig(MariaDBIndexConfig, DBCaseConfig): - M: int | None - ef_search: int | None + M: int | None = None + ef_search: int | None = None index: IndexType = IndexType.HNSW storage_engine: str = "InnoDB" - max_cache_size: int | None + max_cache_size: int | None = None def index_param(self) -> dict: return { diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 9ffbdcece..b04901a95 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -1,4 +1,6 @@ -from pydantic import BaseModel, SecretStr, validator +from typing import Any + +from pydantic import BaseModel, SecretStr, model_validator from ..api import DBCaseConfig, DBConfig, IndexType, MetricType, SQType @@ -19,17 +21,18 @@ def to_dict(self) -> dict: "replica_number": self.replica_number, } - @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"] - ): - return v - if isinstance(v, str | SecretStr) and len(v) == 0: - raise ValueError("Empty string!") - return v + @model_validator(mode="before") + @classmethod + def not_empty_field(cls, data: Any) -> Any: + if isinstance(data, dict): + skip = set(cls.common_short_configs()) | set(cls.common_long_configs()) | {"user", "password"} + for name, v in data.items(): + if name in skip: + continue + if isinstance(v, str) and len(v) == 0: + msg = f"Empty string for field '{name}'!" + raise ValueError(msg) + return data class MilvusIndexConfig(BaseModel): diff --git a/vectordb_bench/backend/clients/oss_opensearch/config.py b/vectordb_bench/backend/clients/oss_opensearch/config.py index 83fed3d58..d153eb429 100644 --- a/vectordb_bench/backend/clients/oss_opensearch/config.py +++ b/vectordb_bench/backend/clients/oss_opensearch/config.py @@ -1,7 +1,8 @@ import logging from enum import Enum +from typing import Any -from pydantic import BaseModel, SecretStr, root_validator, validator +from pydantic import BaseModel, SecretStr, field_validator, model_validator from ..api import DBCaseConfig, DBConfig, MetricType @@ -32,17 +33,18 @@ def to_dict(self) -> dict: "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 + @model_validator(mode="before") + @classmethod + def not_empty_field(cls, data: Any) -> Any: + if isinstance(data, dict): + skip = set(cls.common_short_configs()) | set(cls.common_long_configs()) | {"user", "password", "host"} + for name, v in data.items(): + if name in skip: + continue + if isinstance(v, str) and len(v) == 0: + msg = f"Empty string for field '{name}'!" + raise ValueError(msg) + return data class OSSOS_Engine(Enum): @@ -111,8 +113,9 @@ class OSSOpenSearchIndexConfig(BaseModel, DBCaseConfig): compression_level: str = CompressionLevel.LEVEL_32X oversample_factor: float = 1.0 - @validator("quantization_type", pre=True, always=True) - def validate_quantization_type(cls, value: any): + @field_validator("quantization_type", mode="before") + @classmethod + def validate_quantization_type(cls, value: Any) -> OSSOpenSearchQuantization: """Convert string values to enum""" if not value: return OSSOpenSearchQuantization.NONE @@ -128,19 +131,20 @@ def validate_quantization_type(cls, value: any): return mapping.get(value, OSSOpenSearchQuantization.NONE) - @root_validator - def validate_engine_name(cls, values: dict): + @model_validator(mode="before") + @classmethod + def validate_engine_name(cls, data: Any) -> Any: """Map engine_name string from UI to engine enum""" - if values.get("engine_name"): - engine_name = values["engine_name"].lower() + if isinstance(data, dict) and data.get("engine_name"): + engine_name = data["engine_name"].lower() if engine_name == "faiss": - values["engine"] = OSSOS_Engine.faiss + data["engine"] = OSSOS_Engine.faiss elif engine_name == "lucene": - values["engine"] = OSSOS_Engine.lucene + data["engine"] = OSSOS_Engine.lucene else: log.warning(f"Unknown engine_name: {engine_name}, defaulting to faiss") - values["engine"] = OSSOS_Engine.faiss - return values + data["engine"] = OSSOS_Engine.faiss + return data def __eq__(self, obj: any): return ( diff --git a/vectordb_bench/backend/clients/pgdiskann/config.py b/vectordb_bench/backend/clients/pgdiskann/config.py index 7f83a05c8..8715b1e42 100644 --- a/vectordb_bench/backend/clients/pgdiskann/config.py +++ b/vectordb_bench/backend/clients/pgdiskann/config.py @@ -43,8 +43,8 @@ class PgDiskANNIndexConfig(BaseModel, DBCaseConfig): metric_type: MetricType | None = None create_index_before_load: bool = False create_index_after_load: bool = True - maintenance_work_mem: str | None - max_parallel_workers: int | None + maintenance_work_mem: str | None = None + max_parallel_workers: int | None = None def parse_metric(self) -> str: if self.metric_type == MetricType.L2: @@ -120,10 +120,10 @@ def _optionally_build_set_options( class PgDiskANNImplConfig(PgDiskANNIndexConfig): index: IndexType = IndexType.DISKANN - max_neighbors: int | None - l_value_ib: int | None - pq_param_num_chunks: int | None - l_value_is: float | None + max_neighbors: int | None = None + l_value_ib: int | None = None + pq_param_num_chunks: int | None = None + l_value_is: float | None = None reranking: bool | None = None reranking_metric: str | None = None quantized_fetch_limit: int | None = None diff --git a/vectordb_bench/backend/clients/pgvecto_rs/config.py b/vectordb_bench/backend/clients/pgvecto_rs/config.py index fbb7c5d81..73c2573a0 100644 --- a/vectordb_bench/backend/clients/pgvecto_rs/config.py +++ b/vectordb_bench/backend/clients/pgvecto_rs/config.py @@ -78,7 +78,7 @@ def session_param(self) -> dict[str, str | int]: ... class PgVectoRSHNSWConfig(PgVectoRSIndexConfig): index: IndexType = IndexType.HNSW m: int | None = None - ef_search: int | None + ef_search: int | None = None ef_construction: int | None = None def index_param(self) -> dict[str, str]: @@ -106,8 +106,8 @@ def session_param(self) -> dict[str, str | int]: class PgVectoRSIVFFlatConfig(PgVectoRSIndexConfig): index: IndexType = IndexType.IVFFlat - probes: int | None - lists: int | None + probes: int | None = None + lists: int | None = None def index_param(self) -> dict[str, str]: if self.quantization_type is None: diff --git a/vectordb_bench/backend/clients/pgvector/config.py b/vectordb_bench/backend/clients/pgvector/config.py index 98e82f1c2..7da238a4b 100644 --- a/vectordb_bench/backend/clients/pgvector/config.py +++ b/vectordb_bench/backend/clients/pgvector/config.py @@ -47,8 +47,8 @@ class PgVectorIndexParam(TypedDict): metric: str index_type: str index_creation_with_options: Sequence[dict[str, Any]] - maintenance_work_mem: str | None - max_parallel_workers: int | None + maintenance_work_mem: str | None = None + max_parallel_workers: int | None = None class PgVectorSearchParam(TypedDict): @@ -175,13 +175,13 @@ class PgVectorIVFFlatConfig(PgVectorIndexConfig): a good place to start is sqrt(lists) """ - lists: int | None - probes: int | None + lists: int | None = None + probes: int | None = None index: IndexType = IndexType.ES_IVFFlat maintenance_work_mem: str | None = None max_parallel_workers: int | None = None quantization_type: str | None = None - table_quantization_type: str | None + table_quantization_type: str | None = None reranking: bool | None = None quantized_fetch_limit: int | None = None reranking_metric: str | None = None @@ -226,12 +226,12 @@ class PgVectorHNSWConfig(PgVectorIndexConfig): m: int | None # DETAIL: Valid values are between "2" and "100". ef_construction: int | None # ef_construction must be greater than or equal to 2 * m - ef_search: int | None + ef_search: int | None = None index: IndexType = IndexType.ES_HNSW maintenance_work_mem: str | None = None max_parallel_workers: int | None = None quantization_type: str | None = None - table_quantization_type: str | None + table_quantization_type: str | None = None reranking: bool | None = None quantized_fetch_limit: int | None = None reranking_metric: str | None = None diff --git a/vectordb_bench/backend/clients/pgvectorscale/config.py b/vectordb_bench/backend/clients/pgvectorscale/config.py index e22c45c8d..07750cffb 100644 --- a/vectordb_bench/backend/clients/pgvectorscale/config.py +++ b/vectordb_bench/backend/clients/pgvectorscale/config.py @@ -70,14 +70,14 @@ def session_param(self) -> dict: ... class PgVectorScaleStreamingDiskANNConfig(PgVectorScaleIndexConfig): index: IndexType = IndexType.STREAMING_DISKANN - storage_layout: str | None - num_neighbors: int | None - search_list_size: int | None - max_alpha: float | None - num_dimensions: int | None - num_bits_per_dimension: int | None - query_search_list_size: int | None - query_rescore: int | None + storage_layout: str | None = None + num_neighbors: int | None = None + search_list_size: int | None = None + max_alpha: float | None = None + num_dimensions: int | None = None + num_bits_per_dimension: int | None = None + query_search_list_size: int | None = None + query_rescore: int | None = None def index_param(self) -> dict: return { diff --git a/vectordb_bench/backend/clients/qdrant_cloud/config.py b/vectordb_bench/backend/clients/qdrant_cloud/config.py index b2eeb2ce6..45373754f 100644 --- a/vectordb_bench/backend/clients/qdrant_cloud/config.py +++ b/vectordb_bench/backend/clients/qdrant_cloud/config.py @@ -1,6 +1,6 @@ -from typing import TypeVar +from typing import Any, TypeVar -from pydantic import BaseModel, SecretStr, validator +from pydantic import BaseModel, SecretStr, model_validator from ..api import DBCaseConfig, DBConfig, MetricType @@ -25,11 +25,18 @@ def to_dict(self) -> dict: "url": self.url.get_secret_value(), } - @validator("*") - def not_empty_field(cls, v: any, field: any): - if field.name in ["api_key"]: - return v - return super().not_empty_field(v, field) + @model_validator(mode="before") + @classmethod + def not_empty_field(cls, data: Any) -> Any: + if isinstance(data, dict): + skip = set(cls.common_short_configs()) | set(cls.common_long_configs()) | {"api_key"} + for name, v in data.items(): + if name in skip: + continue + if isinstance(v, str) and len(v) == 0: + msg = f"Empty string for field '{name}'!" + raise ValueError(msg) + return data class QdrantIndexConfig(BaseModel, DBCaseConfig): diff --git a/vectordb_bench/backend/clients/tidb/config.py b/vectordb_bench/backend/clients/tidb/config.py index 71fdbad66..72abb9111 100644 --- a/vectordb_bench/backend/clients/tidb/config.py +++ b/vectordb_bench/backend/clients/tidb/config.py @@ -1,6 +1,6 @@ -from typing import TypedDict +from typing import Any, TypedDict -from pydantic import BaseModel, SecretStr, validator +from pydantic import BaseModel, SecretStr, model_validator from ..api import DBCaseConfig, DBConfig, MetricType @@ -35,13 +35,18 @@ def to_dict(self) -> TiDBConfigDict: "ssl_verify_identity": self.ssl, } - @validator("*") - def not_empty_field(cls, v: any, field: any): - if field.name in ["password", "db_label"]: - return v - if isinstance(v, str | SecretStr) and len(v) == 0: - raise ValueError("Empty string!") - return v + @model_validator(mode="before") + @classmethod + def not_empty_field(cls, data: Any) -> Any: + if isinstance(data, dict): + skip = {"password", "db_label"} + for name, v in data.items(): + if name in skip: + continue + if isinstance(v, str) and len(v) == 0: + msg = f"Empty string for field '{name}'!" + raise ValueError(msg) + return data class TiDBIndexConfig(BaseModel, DBCaseConfig): diff --git a/vectordb_bench/backend/dataset.py b/vectordb_bench/backend/dataset.py index d1de9e328..12808e420 100644 --- a/vectordb_bench/backend/dataset.py +++ b/vectordb_bench/backend/dataset.py @@ -7,12 +7,12 @@ import logging import pathlib from enum import Enum -from typing import Any, NamedTuple +from typing import Any, ClassVar, NamedTuple import pandas as pd import polars as pl from pyarrow.parquet import ParquetFile -from pydantic import PrivateAttr, validator +from pydantic import field_validator from vectordb_bench import config from vectordb_bench.base import BaseModel @@ -38,7 +38,7 @@ class BaseDataset(BaseModel): metric_type: MetricType use_shuffled: bool with_gt: bool = False - _size_label: dict[int, SizeLabel] = PrivateAttr() + _size_label: ClassVar[dict[int, SizeLabel]] is_custom: bool = False with_remote_resource: bool = True # for label filter cases @@ -57,8 +57,9 @@ class BaseDataset(BaseModel): gt_id_field: str = "id" gt_neighbors_field: str = "neighbors_id" - @validator("size") - def verify_size(cls, v: int): + @field_validator("size") + @classmethod + def verify_size(cls, v: int) -> int: if v not in cls._size_label: msg = f"Size {v} not supported for the dataset, expected: {cls._size_label.keys()}" raise ValueError(msg) @@ -102,8 +103,9 @@ class CustomDataset(BaseDataset): scalar_labels_file: str = "scalar_labels.parquet" label_percentages: list[float] = [] - @validator("size") - def verify_size(cls, v: int): + @field_validator("size") + @classmethod + def verify_size(cls, v: int) -> int: return v @property @@ -136,7 +138,7 @@ class LAION(BaseDataset): metric_type: MetricType = MetricType.L2 use_shuffled: bool = False with_gt: bool = True - _size_label: dict = { + _size_label: ClassVar[dict] = { 100_000_000: SizeLabel(100_000_000, "LARGE", 100), } @@ -146,7 +148,7 @@ class GIST(BaseDataset): dim: int = 960 metric_type: MetricType = MetricType.L2 use_shuffled: bool = False - _size_label: dict = { + _size_label: ClassVar[dict] = { 100_000: SizeLabel(100_000, "SMALL", 1), 1_000_000: SizeLabel(1_000_000, "MEDIUM", 1), } @@ -158,7 +160,7 @@ class Cohere(BaseDataset): metric_type: MetricType = MetricType.COSINE use_shuffled: bool = config.USE_SHUFFLED_DATA with_gt: bool = True - _size_label: dict = { + _size_label: ClassVar[dict] = { 100_000: SizeLabel(100_000, "SMALL", 1), 1_000_000: SizeLabel(1_000_000, "MEDIUM", 1), 10_000_000: SizeLabel(10_000_000, "LARGE", 10), @@ -196,7 +198,7 @@ class Bioasq(BaseDataset): metric_type: MetricType = MetricType.COSINE use_shuffled: bool = config.USE_SHUFFLED_DATA with_gt: bool = True - _size_label: dict = { + _size_label: ClassVar[dict] = { 1_000_000: SizeLabel(1_000_000, "MEDIUM", 1), 10_000_000: SizeLabel(10_000_000, "LARGE", 10), } @@ -232,7 +234,7 @@ class Glove(BaseDataset): dim: int = 200 metric_type: MetricType = MetricType.COSINE use_shuffled: bool = False - _size_label: dict = {1_000_000: SizeLabel(1_000_000, "MEDIUM", 1)} + _size_label: ClassVar[dict] = {1_000_000: SizeLabel(1_000_000, "MEDIUM", 1)} class SIFT(BaseDataset): @@ -240,7 +242,7 @@ class SIFT(BaseDataset): dim: int = 128 metric_type: MetricType = MetricType.L2 use_shuffled: bool = False - _size_label: dict = { + _size_label: ClassVar[dict] = { 500_000: SizeLabel( 500_000, "SMALL", @@ -257,7 +259,7 @@ class OpenAI(BaseDataset): metric_type: MetricType = MetricType.COSINE use_shuffled: bool = config.USE_SHUFFLED_DATA with_gt: bool = True - _size_label: dict = { + _size_label: ClassVar[dict] = { 50_000: SizeLabel(50_000, "SMALL", 1), 500_000: SizeLabel(500_000, "MEDIUM", 1), 5_000_000: SizeLabel(5_000_000, "LARGE", 10), diff --git a/vectordb_bench/backend/task_runner.py b/vectordb_bench/backend/task_runner.py index 8224a0415..89a70d5fb 100644 --- a/vectordb_bench/backend/task_runner.py +++ b/vectordb_bench/backend/task_runner.py @@ -73,7 +73,7 @@ def __hash__(self) -> int: ) def display(self) -> dict: - c_dict = self.ca.dict( + c_dict = self.ca.model_dump( include={ "label": True, "name": True, diff --git a/vectordb_bench/base.py b/vectordb_bench/base.py index 502d5fa49..694308d90 100644 --- a/vectordb_bench/base.py +++ b/vectordb_bench/base.py @@ -1,5 +1,6 @@ from pydantic import BaseModel as PydanticBaseModel +from pydantic import ConfigDict -class BaseModel(PydanticBaseModel, arbitrary_types_allowed=True): - pass +class BaseModel(PydanticBaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True, validate_default=True) diff --git a/vectordb_bench/frontend/components/custom/getCustomConfig.py b/vectordb_bench/frontend/components/custom/getCustomConfig.py index a1ddfb737..03831a83c 100644 --- a/vectordb_bench/frontend/components/custom/getCustomConfig.py +++ b/vectordb_bench/frontend/components/custom/getCustomConfig.py @@ -62,14 +62,16 @@ def get_custom_streaming_configs(): def save_custom_configs(custom_configs: list[CustomDatasetConfig]): with open(config.CUSTOM_CONFIG_DIR, "w") as f: - json.dump([custom_config.dict() for custom_config in custom_configs], f, indent=4) + json.dump([custom_config.model_dump() for custom_config in custom_configs], f, indent=4) def save_all_custom_configs( performance_configs: list[CustomCaseConfig], streaming_configs: list[CustomStreamingCaseConfig] ): """Save both performance and streaming configs to the same JSON file""" - all_configs = [config.dict() for config in performance_configs] + [config.dict() for config in streaming_configs] + all_configs = [config.model_dump() for config in performance_configs] + [ + config.model_dump() for config in streaming_configs + ] with open(config.CUSTOM_CONFIG_DIR, "w") as f: json.dump(all_configs, f, indent=4) diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index e8c81c1d1..825f89893 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -119,7 +119,7 @@ def get_custom_case_items() -> list[UICaseItem]: CaseConfig( case_id=CaseType.PerformanceCustomDataset, custom_case={ - **custom_config.dict(), + **custom_config.model_dump(), "use_filter": False, }, ) @@ -140,7 +140,7 @@ def get_custom_case_items() -> list[UICaseItem]: CaseConfig( case_id=CaseType.PerformanceCustomDataset, custom_case={ - **custom_config.dict(), + **custom_config.model_dump(), "use_filter": True, "label_percentage": label_percentage, }, @@ -174,7 +174,7 @@ def get_custom_streaming_case_items() -> list[UICaseItem]: case_id=CaseType.StreamingCustomDataset, custom_case={ "description": custom_config.description, - "dataset_config": custom_config.dataset_config.dict(), + "dataset_config": custom_config.dataset_config.model_dump(), }, ) ], diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 27be85b6e..88b162ae3 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -200,7 +200,7 @@ def k(self, value): ''' def __hash__(self) -> int: - return hash(self.json()) + return hash(self.model_dump_json()) @property def case(self) -> Case: @@ -308,7 +308,7 @@ def write_db_file(self, result_dir: pathlib.Path, partial: Self, db: str): log.info(f"write results to disk {result_file}") with pathlib.Path(result_file).open("w") as f: - b = partial.json(exclude={"db_config": {"password", "api_key"}}) + b = partial.model_dump_json(exclude={"db_config": {"password", "api_key"}}) f.write(b) def get_case_config(case_config: CaseConfig) -> dict[CaseConfig]: diff --git a/vectordb_bench/restful/format_res.py b/vectordb_bench/restful/format_res.py index 2e289ec3b..326986319 100644 --- a/vectordb_bench/restful/format_res.py +++ b/vectordb_bench/restful/format_res.py @@ -63,7 +63,7 @@ def format_results(test_results: list[TestResult], task_label: str) -> list[dict db_label=task_config.db_config.db_label, version=task_config.db_config.version, note=task_config.db_config.note, - params=task_config.db_case_config.dict(), + params=task_config.db_case_config.model_dump(), case_name=case.name, dataset=dataset.full_name, dim=dataset.dim, @@ -71,6 +71,6 @@ def format_results(test_results: list[TestResult], task_label: str) -> list[dict filter_rate=filter_.filter_rate, k=task_config.case_config.k, **metrics, - ).dict() + ).model_dump() ) return results