diff --git a/Makefile b/Makefile deleted file mode 100644 index 8a899c7b1..000000000 --- a/Makefile +++ /dev/null @@ -1,88 +0,0 @@ -PYTHON ?= /home/rowan/.local/share/uv/tools/pip/bin/python -BENCH = $(PYTHON) -m vectordb_bench.cli.vectordbbench -COMMON = --drop-old --skip-search-concurrent - -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: - $(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: - $(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 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/vectordb_bench/backend/clients/antfly/antfly.py b/vectordb_bench/backend/clients/antfly/antfly.py index 057c6397a..77176e4fd 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") @@ -35,50 +38,68 @@ def __init__( self.collection_name = collection_name self.dim = dim - base_url = f"http://{db_config['host']}:{db_config['port']}/api/v1" - self._base_url = base_url + self._metadata_base_url = f"http://{db_config['host']}:{db_config['port']}/api/v1" + self._store_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, - "field": SOURCE_FIELD, - **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: @@ -87,7 +108,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"}, @@ -106,6 +126,29 @@ 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: + raise ValueError( + f"Antfly direct store search currently requires exactly one shard; found {len(shards)} shards" + ) + self._direct_shard_id = next(iter(shards)) + def _index_status_is_ready( self, payload: dict | None, @@ -151,17 +194,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: @@ -177,11 +229,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) @@ -189,7 +303,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( @@ -209,27 +323,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}") @@ -244,33 +346,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 {}