From c804dfd902285ce2d7c0fe3a3d9832a283ba1c11 Mon Sep 17 00:00:00 2001 From: Thor Whalen <1906276+thorwhalen@users.noreply.github.com> Date: Fri, 22 May 2026 09:54:42 +0200 Subject: [PATCH] Live-verify server/managed backend adapters; fix 8 bugs found Exercised the 9 previously-unverified backend adapters (issue #11) against real engines: Docker containers for pgvector / redis / elasticsearch / weaviate / mongodb, the embedded Milvus Lite engine for milvus, and a loadable-extension Python for sqlite_vec. All 344 contract tests now pass across 13 backends. Adapter bugs fixed: - sqlite_vec: vec0 tables reject INSERT OR REPLACE (delete+insert instead); a KNN query through a JOIN drops the required `k` constraint (subquery form); an over-broad `except OperationalError` silently swallowed query errors; a NULL distance for a zero-norm vector crashed scoring. - milvus: used tempfile.mkstemp (creates a file) but milvus-lite >=3.0 treats the path as a directory -> FileExistsError. - pgvector: the vector(N) type modifier cannot be a bound parameter; executemany is a cursor (not connection) method in psycopg 3; the query vector needs an explicit ::vector cast. - elasticsearch: lazily-created collections were invisible to list / contains / get (the index is now created eagerly, the dense_vector field added on first write); _keys sorted on the disallowed _id field (now uses helpers.scan). - weaviate: collections.get() never raises in v4, so _class_exists always reported True -> create / get / delete were broken (now collections.exists()). - mongodb: _raw_to_document was undefined (NameError in _read); the Atlas vector search index is now created and awaited automatically on the first search, and metadata filtering moved client-side (Atlas can only pre-filter on index-declared fields). Test harness: - tests/docker-compose.yml: one container per server backend. - tests/conftest.py: sweeps embedded + server backends; server backends are TCP-probed and skipped when unreachable, so CI stays green without them. - tests/README.md: how to run the server-backend suite. Also: de-degenerate three test inputs that used all-zero vectors (undefined cosine similarity, which Elasticsearch rejects); add milvus-lite to the milvus extra so the embedded engine works out of the box. Refs #11 --- pyproject.toml | 2 +- tests/README.md | 60 ++++++ tests/conftest.py | 180 ++++++++++++++++-- tests/docker-compose.yml | 82 ++++++++ tests/test_core.py | 13 +- vd/backends/elasticsearch.py | 106 ++++++----- vd/backends/milvus.py | 8 +- vd/backends/mongodb.py | 358 +++++++++++++++++------------------ vd/backends/pgvector.py | 32 ++-- vd/backends/sqlite_vec.py | 61 ++++-- vd/backends/weaviate.py | 10 +- 11 files changed, 622 insertions(+), 290 deletions(-) create mode 100644 tests/README.md create mode 100644 tests/docker-compose.yml diff --git a/pyproject.toml b/pyproject.toml index 05ca3c4..7cc8005 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,7 +34,7 @@ duckdb = ["duckdb"] pgvector = ["pgvector", "psycopg[binary]"] pinecone = ["pinecone"] weaviate = ["weaviate-client"] -milvus = ["pymilvus"] +milvus = ["pymilvus", "milvus-lite ; platform_system != 'Windows'"] redis = ["redis", "numpy"] elasticsearch = ["elasticsearch"] mongodb = ["pymongo"] diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..9d6edf4 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,60 @@ +# vd test suite + +The headline file is `test_core.py` — the **facade contract suite**. Every test +there takes the `client` fixture, which is parametrized over every backend, so a +single test asserts the same behaviour across all of them. That is how `vd` +proves one uniform interface over many vector databases. + +## Backends the suite sweeps over + +`conftest.py` defines two groups: + +- **Embedded backends** — `memory`, `chroma`, `faiss`, `duckdb`, `lancedb`, + `qdrant`, `sqlite_vec`, `milvus`. No server needed; each test gets a fresh + client. `sqlite_vec` is skipped on a Python whose `sqlite3` lacks + loadable-extension support; `milvus` runs against the embedded **Milvus Lite** + engine and is skipped if `milvus-lite` is not installed. +- **Server backends** — `pgvector`, `redis`, `elasticsearch`, `weaviate`, + `mongodb`. Each needs a running container. A backend whose port is not open is + **skipped**, so the suite stays green in a plain CI environment. + +## Running against the server backends + +```bash +# Start every server backend (first run pulls images — a few minutes). +docker compose -f tests/docker-compose.yml up -d + +# Wait until all five report (healthy), then run the suite. +docker compose -f tests/docker-compose.yml ps + +pytest tests/ + +# Tear down and wipe volumes when done. +docker compose -f tests/docker-compose.yml down -v +``` + +Run one backend at a time with `-k`: + +```bash +pytest tests/test_core.py -k pgvector +``` + +## Connection settings + +The server backends use these defaults; override with environment variables: + +| Backend | Env var | Default | +|-----------------|-----------------------|------------------------------------------------------| +| `pgvector` | `VD_PGVECTOR_DSN` | `postgresql://vd:vd@localhost:5432/vd` | +| `redis` | `VD_REDIS_HOST` / `VD_REDIS_PORT` | `localhost` / `6379` | +| `elasticsearch` | `VD_ELASTICSEARCH_URL`| `http://localhost:9200` | +| `weaviate` | `VD_WEAVIATE_HOST` | `localhost` | +| `mongodb` | `VD_MONGODB_URI` | `mongodb://localhost:27018/?directConnection=true` | + +`mongodb` maps to host port **27018** (not 27017) to avoid colliding with a +developer's native `mongod`. It uses the `mongodb-atlas-local` image because +`$vectorSearch` requires an Atlas-capable deployment. + +`milvus` needs no container — it is verified against the embedded Milvus Lite +engine, which exercises the same adapter code path as a Milvus server (only the +client constructor differs). diff --git a/tests/conftest.py b/tests/conftest.py index 06812e5..b4a09df 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,13 +1,32 @@ """ Pytest fixtures shared across the vd test suite. -The headline fixture is :func:`client` — parametrized over every backend that -can actually run in a plain CI environment (no servers, no cloud accounts). -A test that takes ``client`` runs once per backend, which is how the test -suite proves the facade contract holds uniformly. +The headline fixture is :func:`client` — parametrized over every backend vd +can reach in the current environment. A test that takes ``client`` runs once +per backend, which is how the suite proves the facade contract holds uniformly. + +Two kinds of backend run: + +- **Embedded backends** (``memory``, ``chroma``, ``faiss``, ``duckdb``, + ``lancedb``, ``qdrant``, ``sqlite_vec``, ``milvus``) need no server. Each + test gets a fresh client. ``sqlite_vec`` is skipped on a Python whose + ``sqlite3`` lacks loadable-extension support; ``milvus`` runs against the + embedded Milvus Lite engine and is skipped if ``milvus-lite`` is absent. +- **Server backends** (``pgvector``, ``redis``, ``elasticsearch``, + ``weaviate``, ``mongodb``) need a running container — see + ``tests/docker-compose.yml``. Each is TCP-probed and **skipped** when its + container is down, so the suite stays green in a plain CI environment. + +Connection settings for the server backends are environment-overridable +(``VD_PGVECTOR_DSN``, ``VD_REDIS_HOST``/``VD_REDIS_PORT``, +``VD_ELASTICSEARCH_URL``, ``VD_WEAVIATE_HOST``, ``VD_MONGODB_URI``). """ import hashlib +import importlib.util +import os +import socket +import sqlite3 import pytest @@ -16,10 +35,135 @@ #: Embedding dimension used by the test embedder. EMBED_DIM = 16 -#: Backends exercised end-to-end here: embedded / pip-only, no server needed. -#: (sqlite_vec is intentionally excluded — some Python builds ship a sqlite3 -#: with extension loading disabled, which the backend needs.) -TESTABLE_BACKENDS = ["memory", "chroma", "faiss", "duckdb", "lancedb", "qdrant"] +#: Backends that need no server. Each test gets a fresh client. +EMBEDDED_BACKENDS = [ + "memory", + "chroma", + "faiss", + "duckdb", + "lancedb", + "qdrant", + "sqlite_vec", + "milvus", # verified against the embedded Milvus Lite engine +] + +#: Server backends — each needs a container (``tests/docker-compose.yml``). +#: ``probe`` is TCP-probed; the backend is skipped when the port is closed. +#: ``connect_kwargs`` builds the :func:`vd.connect` arguments (env-overridable). +SERVER_BACKENDS = { + "pgvector": { + "probe": ("localhost", 5432), + "connect_kwargs": lambda: { + "dsn": os.environ.get( + "VD_PGVECTOR_DSN", "postgresql://vd:vd@localhost:5432/vd" + ) + }, + }, + "redis": { + "probe": ("localhost", 6379), + "connect_kwargs": lambda: { + "host": os.environ.get("VD_REDIS_HOST", "localhost"), + "port": int(os.environ.get("VD_REDIS_PORT", "6379")), + }, + }, + "elasticsearch": { + "probe": ("localhost", 9200), + "connect_kwargs": lambda: { + "url": os.environ.get("VD_ELASTICSEARCH_URL", "http://localhost:9200") + }, + }, + "weaviate": { + "probe": ("localhost", 8080), + "connect_kwargs": lambda: { + "host": os.environ.get("VD_WEAVIATE_HOST", "localhost") + }, + }, + "mongodb": { + # Host port 27018 — see tests/docker-compose.yml (avoids colliding + # with a developer's native mongod on the default 27017). + "probe": ("localhost", 27018), + "connect_kwargs": lambda: { + "uri": os.environ.get( + "VD_MONGODB_URI", "mongodb://localhost:27018/?directConnection=true" + ) + }, + }, +} + +#: Every backend the parametrized ``client`` fixture sweeps over. +ALL_BACKENDS = EMBEDDED_BACKENDS + list(SERVER_BACKENDS) + + +# --------------------------------------------------------------------------- # +# Availability probes +# --------------------------------------------------------------------------- # + + +def _tcp_open(host: str, port: int, timeout: float = 0.5) -> bool: + """Return ``True`` if a TCP connection to ``host:port`` succeeds.""" + try: + with socket.create_connection((host, port), timeout=timeout): + return True + except OSError: + return False + + +def _sqlite_ext_supported() -> bool: + """Return ``True`` if this Python's sqlite3 supports loadable extensions.""" + try: + conn = sqlite3.connect(":memory:") + ok = hasattr(conn, "enable_load_extension") + conn.close() + return ok + except Exception: + return False + + +def _unavailable_reason(name: str) -> str | None: + """Return a skip reason for backend ``name``, or ``None`` if it can run.""" + if name == "sqlite_vec" and not _sqlite_ext_supported(): + return "sqlite3 was built without loadable-extension support" + if name == "milvus" and importlib.util.find_spec("milvus_lite") is None: + return "milvus-lite not installed (embedded Milvus engine unavailable)" + if name in SERVER_BACKENDS: + host, port = SERVER_BACKENDS[name]["probe"] + if not _tcp_open(host, port): + return ( + f"{name!r} server unreachable at {host}:{port} " + f"— start it with tests/docker-compose.yml" + ) + return None + + +def _connect_kwargs(name: str) -> dict: + """Return the :func:`vd.connect` kwargs for backend ``name``.""" + if name in SERVER_BACKENDS: + return SERVER_BACKENDS[name]["connect_kwargs"]() + return {} + + +def _drop_all_collections(client) -> None: + """ + Delete every collection on ``client`` — best-effort. + + Server backends keep state across runs; embedded backends are fresh each + time. Dropping everything before and after each test makes a run against a + live server idempotent (a re-run does not trip "already exists"). + """ + try: + names = list(client.list_collections()) + except Exception: + return + for name in names: + try: + client.delete_collection(name) + except Exception: + pass + + +# --------------------------------------------------------------------------- # +# Test embedder +# --------------------------------------------------------------------------- # def make_embedder(): @@ -41,20 +185,32 @@ def embedder(): return make_embedder() -@pytest.fixture(params=TESTABLE_BACKENDS) +# --------------------------------------------------------------------------- # +# Backend fixtures +# --------------------------------------------------------------------------- # + + +@pytest.fixture(params=ALL_BACKENDS) def backend_name(request): - """Each installed testable backend, one at a time.""" + """Each reachable backend, one at a time; unreachable ones are skipped.""" name = request.param if name not in vd.list_backends(): pytest.skip(f"backend {name!r} is not installed") + reason = _unavailable_reason(name) + if reason: + pytest.skip(reason) return name @pytest.fixture def client(backend_name, embedder): - """A fresh, connected client for each testable backend (with an embedder).""" - connection = vd.connect(backend_name, embedder=embedder) + """A fresh, connected client for each backend (with an embedder).""" + connection = vd.connect( + backend_name, embedder=embedder, **_connect_kwargs(backend_name) + ) + _drop_all_collections(connection) yield connection + _drop_all_collections(connection) if hasattr(connection, "close"): try: connection.close() diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml new file mode 100644 index 0000000..f51898d --- /dev/null +++ b/tests/docker-compose.yml @@ -0,0 +1,82 @@ +# Server backends for vd's live-verification test suite. +# +# Brings up one container per server-backed vd backend. The pytest suite +# (see conftest.py) TCP-probes each port and runs the contract tests against +# whatever is reachable; backends whose container is down are skipped, so the +# suite stays green in a plain CI environment. +# +# docker compose -f tests/docker-compose.yml up -d # start +# pytest tests/ # run the suite +# docker compose -f tests/docker-compose.yml down -v # stop + wipe +# +# milvus is intentionally absent: the milvus adapter is verified against the +# embedded Milvus Lite engine (no server), which exercises the same code path. + +services: + pgvector: + image: pgvector/pgvector:pg17 + environment: + POSTGRES_USER: vd + POSTGRES_PASSWORD: vd + POSTGRES_DB: vd + ports: + - "5432:5432" + healthcheck: + test: ["CMD-SHELL", "pg_isready -U vd"] + interval: 3s + timeout: 5s + retries: 20 + + redis: + image: redis/redis-stack-server:latest + ports: + - "6379:6379" + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 3s + timeout: 5s + retries: 20 + + elasticsearch: + image: docker.elastic.co/elasticsearch/elasticsearch:9.0.3 + environment: + discovery.type: single-node + xpack.security.enabled: "false" + ES_JAVA_OPTS: "-Xms512m -Xmx512m" + ports: + - "9200:9200" + healthcheck: + test: ["CMD-SHELL", "curl -sf http://localhost:9200/_cluster/health || exit 1"] + interval: 5s + timeout: 5s + retries: 30 + + weaviate: + image: cr.weaviate.io/semitechnologies/weaviate:1.28.4 + command: ["--host", "0.0.0.0", "--port", "8080", "--scheme", "http"] + environment: + AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED: "true" + PERSISTENCE_DATA_PATH: "/var/lib/weaviate" + DEFAULT_VECTORIZER_MODULE: "none" + ENABLE_MODULES: "" + QUERY_DEFAULTS_LIMIT: "25" + ports: + - "8080:8080" + - "50051:50051" + healthcheck: + test: ["CMD-SHELL", "wget -qO- http://localhost:8080/v1/.well-known/ready || exit 1"] + interval: 5s + timeout: 5s + retries: 30 + + mongodb: + image: mongodb/mongodb-atlas-local:latest + # Host port 27018 (not 27017) to avoid colliding with a developer's + # native mongod. The container still listens on 27017 internally. + ports: + - "27018:27017" + healthcheck: + test: ["CMD-SHELL", "mongosh --quiet --eval 'db.runCommand({ ping: 1 })' || exit 1"] + interval: 5s + timeout: 5s + retries: 30 diff --git a/tests/test_core.py b/tests/test_core.py index eae376d..b75433b 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -166,7 +166,9 @@ def test_search_with_comparison_filter(client): col = client.create_collection("search_cmp") for i in range(6): v = [0.0] * 16 - v[0] = i / 6 + # (i + 1), not i — a zero-magnitude vector has undefined cosine + # similarity and some backends (Elasticsearch) reject it outright. + v[0] = (i + 1) / 6 col[str(i)] = _doc(str(i), v, n=i) hits = list(col.search([0.5] * 16, limit=10, filter={"n": {"$gte": 3}})) assert sorted(int(h["id"]) for h in hits) == [3, 4, 5] @@ -176,7 +178,8 @@ def test_search_with_in_filter(client): col = client.create_collection("search_in") for i in range(5): v = [0.0] * 16 - v[0] = i / 5 + # (i + 1), not i — avoid a zero-magnitude vector (see comment above). + v[0] = (i + 1) / 5 col[str(i)] = _doc(str(i), v, n=i) hits = list(col.search([0.5] * 16, limit=10, filter={"n": {"$in": [1, 3]}})) assert sorted(int(h["id"]) for h in hits) == [1, 3] @@ -229,6 +232,8 @@ def test_upsert_is_idempotent(client): def test_dimension_mismatch_raises(client): col = client.create_collection("dimcheck") - col["ok"] = Document(id="ok", text="t", vector=[0.0] * 16) + # A non-zero vector — a zero-magnitude one has undefined cosine similarity + # and some backends (Elasticsearch) reject it outright. + col["ok"] = Document(id="ok", text="t", vector=[1.0] * 16) with pytest.raises(ValueError): - col["bad"] = Document(id="bad", text="t", vector=[0.0] * 8) + col["bad"] = Document(id="bad", text="t", vector=[1.0] * 8) diff --git a/vd/backends/elasticsearch.py b/vd/backends/elasticsearch.py index 1ab904b..7de6575 100644 --- a/vd/backends/elasticsearch.py +++ b/vd/backends/elasticsearch.py @@ -195,16 +195,28 @@ def native(self) -> Elasticsearch: def _ensure_index(self) -> None: """ - Create the Elasticsearch index on the first write if it does not exist. - - The dimension must be known (set) before this is called, which is - guaranteed because :meth:`~vd.base.AbstractCollection._vet_vector` - sets ``self.dimension`` from the first vector before ``_write`` runs. + Ensure the ``embedding`` dense_vector field is mapped before a write. + + The index itself is created eagerly by + :meth:`ElasticsearchClient.create_collection`, so an empty collection + is still visible to ``list_collections`` / ``in`` / ``get_collection``. + The ``dense_vector`` field, however, needs the vector dimension — not + known until the first vector arrives — so it is added here, lazily, via + ``put_mapping``. ``self.dimension`` is guaranteed set by then because + :meth:`~vd.base.AbstractCollection._vet_vector` runs before ``_write``. """ + full_mapping = _build_mapping(self.dimension, self.metric) if not _index_exists(self._es, self.name): - self._es.indices.create( + # Index missing (collection created outside the adapter, or the + # index was dropped) — create it whole. + self._es.indices.create(index=self.name, mappings=full_mapping) + return + mapping = self._es.indices.get_mapping(index=self.name) + properties = mapping[self.name]["mappings"].get("properties", {}) + if "embedding" not in properties: + self._es.indices.put_mapping( index=self.name, - mappings=_build_mapping(self.dimension, self.metric), + properties={"embedding": full_mapping["properties"]["embedding"]}, ) # ----- raw primitives --------------------------------------------------- # @@ -276,37 +288,26 @@ def _keys(self) -> Iterator[str]: """ Iterate all document ids in the index. - Uses a ``match_all`` query with pagination (``search_after``) to - retrieve ids page by page. For very large indices (tens of millions of - documents) prefer the Scroll API or the ``helpers.scan`` utility; the - ``search_after`` approach used here avoids the deprecated Scroll deep - pagination but still loads all ids into memory. + Uses the scroll-based ``helpers.scan`` to walk every document. + Elasticsearch disallows sorting (fielddata access) on the ``_id`` + field, which rules out an ``_id``-keyed ``search_after`` pagination; + ``scan`` needs no sort. Ids are materialised into a list (as the + previous ``search_after`` implementation also did). """ if not _index_exists(self._es, self.name): return iter(()) + from elasticsearch.helpers import scan - ids: list[str] = [] - search_after: Any = None - while True: - body: dict = { - "query": {"match_all": {}}, - "size": _KEYS_PAGE_SIZE, - "_source": False, - "sort": [{"_id": "asc"}], - } - if search_after is not None: - body["search_after"] = search_after - - response = self._es.search(index=self.name, body=body) - hits = response["hits"]["hits"] - if not hits: - break - for hit in hits: - ids.append(hit["_id"]) - search_after = hits[-1]["sort"] - if len(hits) < _KEYS_PAGE_SIZE: - break - + ids = [ + hit["_id"] + for hit in scan( + self._es, + index=self.name, + query={"query": {"match_all": {}}}, + _source=False, + size=_KEYS_PAGE_SIZE, + ) + ] return iter(ids) def _count(self) -> int: @@ -441,8 +442,12 @@ def create_collection( """ Create a new collection (Elasticsearch index). - The index is created **lazily** on the first write when ``dimension`` - is ``None``, or **eagerly** here when ``dimension`` is supplied. + The index is created **eagerly** so the collection is immediately + visible to :meth:`list_collections`, ``in``, and :meth:`get_collection` + even before its first write. The ``embedding`` ``dense_vector`` field + is added with a full mapping now when ``dimension`` is supplied, or + **lazily** on the first write otherwise (Elasticsearch needs the vector + dimension to map a ``dense_vector`` field). Parameters ---------- @@ -450,8 +455,8 @@ def create_collection( Index name. **Must be lowercase** — Elasticsearch rejects names with uppercase letters. dimension : int, optional - Vector dimension. Required for eager index creation; may be - deferred until the first write. + Vector dimension. When given, the ``dense_vector`` field is mapped + now; otherwise it is deferred until the first write. metric : str Distance metric: ``"cosine"``, ``"dot"``, or ``"l2"``. **index_config @@ -472,16 +477,23 @@ def create_collection( dimension=dimension, metric=metric, ) - if dimension is not None: - # Eager index creation: build the mapping now. - settings = index_config if index_config else None - create_kwargs: dict[str, Any] = { - "index": name, - "mappings": _build_mapping(dimension, metric), + # Always create the index now (an empty collection must still be + # listable). The dense_vector field needs the dimension, so it is only + # mapped here when known — otherwise _ensure_index adds it on first write. + mappings: dict[str, Any] = { + "properties": { + "text": {"type": "text"}, + "metadata": {"type": "object", "enabled": False}, } - if settings: - create_kwargs["settings"] = settings - self._client.indices.create(**create_kwargs) + } + if dimension is not None: + mappings["properties"]["embedding"] = _build_mapping(dimension, metric)[ + "properties" + ]["embedding"] + create_kwargs: dict[str, Any] = {"index": name, "mappings": mappings} + if index_config: + create_kwargs["settings"] = index_config + self._client.indices.create(**create_kwargs) return col def get_collection(self, name: str) -> ElasticsearchCollection: diff --git a/vd/backends/milvus.py b/vd/backends/milvus.py index a4337e9..54eae82 100644 --- a/vd/backends/milvus.py +++ b/vd/backends/milvus.py @@ -45,6 +45,7 @@ from __future__ import annotations +import os import tempfile from typing import Any, Callable, Iterable, Iterator, Optional @@ -514,8 +515,11 @@ def __init__( elif uri is not None: self._client = MilvusClient(uri=uri, token=token) else: - # Milvus Lite: create a temporary .db file. - _fd, tmp_path = tempfile.mkstemp(suffix=".db", prefix="vd_milvus_") + # Milvus Lite: an embedded, server-less engine. Hand it a fresh + # ".db" path inside a new temp directory. milvus-lite >=3.0 treats + # the path as a data directory it creates itself, so it must not + # already exist as a file (which ``tempfile.mkstemp`` would create). + tmp_path = os.path.join(tempfile.mkdtemp(prefix="vd_milvus_"), "milvus.db") self._client = MilvusClient(tmp_path) #: Track metrics for collections created via this client so they survive diff --git a/vd/backends/mongodb.py b/vd/backends/mongodb.py index 0013ab5..e07fbcc 100644 --- a/vd/backends/mongodb.py +++ b/vd/backends/mongodb.py @@ -29,28 +29,17 @@ pipeline. The score returned by ``$meta: "vectorSearchScore"`` is *higher-is-better* and is used directly as the ``vd`` result ``score``. -**IMPORTANT — the Atlas vector search index must be created out-of-band.** -``MongoDBClient.create_collection`` creates the *MongoDB collection* (the -document container) but it **cannot** create the Atlas vector search index. -The vector search index must be created separately, before ``search()`` is -called, via: - -1. The Atlas UI: *Atlas → Browse Collections → Search Indexes → Create - Search Index → Atlas Vector Search*, or -2. The Atlas Data API / Admin API, or -3. The MongoDB Atlas Terraform provider. - -When you create the index, set: - -- *field*: ``"embedding"`` -- *type*: ``"knnVector"`` -- *dimensions*: the embedding dimension you are using (e.g. ``1536``) -- *similarity*: ``"cosine"`` (or ``"euclidean"`` / ``"dotProduct"`` to match - your ``metric`` choice) - -The default index name expected by this adapter is ``"vector_index"``; pass -``vector_index="my_name"`` to ``MongoDBClient`` or ``index="my_name"`` to -``collection.search()`` to override. +**The Atlas vector search index is created automatically.** On the first +:meth:`~MongoDBCollection.search`, this adapter creates the vector search +index — on the ``"embedding"`` field, with the collection's dimension and +metric — via ``create_search_index``, then blocks until Atlas reports it +queryable. This makes the backend behave like every other ``vd`` adapter +(which all create their own index). Index creation requires an Atlas-capable +deployment: Atlas, or a local ``mongodb-atlas-local`` / AtlasCLI deployment — +it is **not** available on a plain ``mongod``. The index name defaults to +``"vector_index"``; override it with ``vector_index="my_name"`` on +:class:`MongoDBClient`, or pass ``index="my_name"`` to +``collection.search()`` to use a different, externally-managed index. On the **M0 free tier**, only one vector search index is allowed per cluster. @@ -65,13 +54,14 @@ try: from pymongo import MongoClient from pymongo.collection import Collection as PymongoCollection - from pymongo.errors import CollectionInvalid + from pymongo.errors import CollectionInvalid, OperationFailure except ImportError as e: # pragma: no cover raise ImportError( "The mongodb backend needs the 'pymongo' package. " "Install with: pip install pymongo" ) from e +from vd.backends._helpers import apply_client_filter, overfetch_limit from vd.base import ( AbstractClient, AbstractCollection, @@ -105,126 +95,6 @@ } -# --------------------------------------------------------------------------- -# Filter translation helpers -# --------------------------------------------------------------------------- - - -def _prefix_field(key: str) -> str: - """ - Return the MongoDB document path for a vd filter field key. - - User metadata is stored under the ``metadata`` sub-document, so every - plain field key is prefixed with ``"metadata."``. Operator keys that - start with ``"$"`` are returned unchanged. - - Examples - -------- - >>> _prefix_field("year") - 'metadata.year' - >>> _prefix_field("$and") - '$and' - """ - return key if key.startswith("$") else f"metadata.{key}" - - -def _to_mongo_filter(ast: Optional[Filter]) -> Optional[dict]: - """ - Translate a canonical ``vd`` filter AST to a MongoDB Query Language document. - - The canonical vd filter dialect is already modelled after MQL, so the - translation is nearly identity. The one structural change: every - *field* key (not operator key) must be prefixed with ``"metadata."`` - because user metadata is stored nested under that sub-document in the - MongoDB document schema. Operator keys (``$and``, ``$or``, ``$not``, - ``$eq``, ``$gt``, ...) are left unchanged and MongoDB handles them - natively. - - Logical operators are recursed through transparently. ``$not`` wraps - its sub-document in MQL's ``{"$not": {...}}`` form for field conditions. - - Parameters - ---------- - ast : dict or None - A filter in the canonical ``vd`` dialect (see :mod:`vd.filters`). - ``None`` or empty returns ``None`` (no filter applied). - - Returns - ------- - dict or None - A MongoDB Query Language document ready to pass as ``filter`` in a - ``$vectorSearch`` stage, or ``None`` when no filtering is needed. - - Examples - -------- - >>> _to_mongo_filter(None) - >>> _to_mongo_filter({}) - >>> _to_mongo_filter({"year": 2024}) - {'metadata.year': 2024} - >>> _to_mongo_filter({"year": {"$gte": 2020}}) - {'metadata.year': {'$gte': 2020}} - >>> _to_mongo_filter({"$and": [{"year": 2024}, {"tag": "ai"}]}) - {'$and': [{'metadata.year': 2024}, {'metadata.tag': 'ai'}]} - """ - if not ast: - return None - - result: dict = {} - for key, value in ast.items(): - if key == "$and": - # $and takes a list of sub-filter documents - result["$and"] = [_to_mongo_filter(sub) for sub in value] - elif key == "$or": - # $or takes a list of sub-filter documents - result["$or"] = [_to_mongo_filter(sub) for sub in value] - elif key == "$not": - # $not in the vd AST wraps a single sub-filter; translate each - # field condition inside it with prefixed keys. - result.update(_to_mongo_not(value)) - else: - # Plain field condition — prefix the key with "metadata." - result[_prefix_field(key)] = value - - return result or None - - -def _to_mongo_not(sub: Filter) -> dict: - """ - Translate a ``$not`` sub-filter into MQL negation form. - - MQL ``$not`` operates on a single field condition - (``{"field": {"$not": {...}}}``), whereas the vd ``$not`` wraps an - entire sub-filter. We translate by prefixing each field and wrapping - its condition in ``{"$not": condition}``. Nested logical operators - inside ``$not`` are passed through recursively. - - Parameters - ---------- - sub : dict - The sub-filter expression to negate. - - Returns - ------- - dict - An MQL document expressing the negation. - """ - result: dict = {} - for key, value in sub.items(): - if key in ("$and", "$or", "$not"): - # Logical operators inside $not: recurse via _to_mongo_filter - translated = _to_mongo_filter({key: value}) - if translated: - result.update(translated) - else: - prefixed = _prefix_field(key) - if isinstance(value, dict): - result[prefixed] = {"$not": value} - else: - # Bare equality — express as $not: {$eq: value} - result[prefixed] = {"$not": {"$eq": value}} - return result - - # --------------------------------------------------------------------------- # Helper: build a $vectorSearch aggregation pipeline # --------------------------------------------------------------------------- @@ -235,7 +105,6 @@ def _build_vector_search_pipeline( *, limit: int, index: str, - mongo_filter: Optional[dict], ) -> list[dict]: """ Build the MongoDB aggregation pipeline for a ``$vectorSearch`` query. @@ -247,22 +116,26 @@ def _build_vector_search_pipeline( ``vectorSearchScore`` meta-field (higher-is-better cosine/dot/ euclidean similarity score assigned by Atlas). - The ``numCandidates`` parameter controls the pre-filter candidate pool. - Atlas requires ``numCandidates >= limit``; we use - ``max(limit * _CANDIDATES_MULTIPLIER, _MIN_NUM_CANDIDATES)`` as the - default, matching common practice from the Atlas documentation. + Metadata filtering is **not** done here. Atlas ``$vectorSearch`` can only + pre-filter on fields explicitly declared as ``filter`` fields in the index + definition — but ``vd``'s filter language is open-ended, so filtering is + applied client-side (over-fetch, then :func:`apply_client_filter`), exactly + as the pgvector / redis / weaviate / milvus adapters do. This keeps filter + semantics identical across every backend. + + The ``numCandidates`` parameter controls the candidate pool. Atlas requires + ``numCandidates >= limit``; we use + ``max(limit * _CANDIDATES_MULTIPLIER, _MIN_NUM_CANDIDATES)``. Parameters ---------- vector : list[float] The query embedding vector. limit : int - Number of nearest neighbours to return. + Number of nearest neighbours to return (already over-fetched by the + caller when a filter is present). index : str - Name of the Atlas vector search index (must exist out-of-band). - mongo_filter : dict or None - An MQL filter document (already translated by ``_to_mongo_filter``), - or ``None`` to skip metadata filtering. + Name of the Atlas vector search index. Returns ------- @@ -270,18 +143,16 @@ def _build_vector_search_pipeline( A two-stage aggregation pipeline. """ num_candidates = max(limit * _CANDIDATES_MULTIPLIER, _MIN_NUM_CANDIDATES) - vector_search_stage: dict[str, Any] = { - "index": index, - "path": "embedding", - "queryVector": vector, - "numCandidates": num_candidates, - "limit": limit, - } - if mongo_filter: - vector_search_stage["filter"] = mongo_filter - return [ - {"$vectorSearch": vector_search_stage}, + { + "$vectorSearch": { + "index": index, + "path": "embedding", + "queryVector": vector, + "numCandidates": num_candidates, + "limit": limit, + } + }, { "$project": { "_id": 1, @@ -293,6 +164,36 @@ def _build_vector_search_pipeline( ] +# --------------------------------------------------------------------------- +# Helper: rebuild a vd Document from a raw MongoDB document +# --------------------------------------------------------------------------- + + +def _raw_to_document(raw: dict) -> Document: + """ + Rebuild a :class:`~vd.base.Document` from a raw MongoDB document. + + The stored shape is ``{"_id": id, "text": text, "embedding": vector, + "metadata": {...}}`` (see :meth:`MongoDBCollection._write`); this is the + inverse mapping. + + Parameters + ---------- + raw : dict + A document as returned by ``pymongo``'s ``find_one`` / ``aggregate``. + + Returns + ------- + Document + """ + return Document( + id=raw["_id"], + text=raw.get("text", ""), + vector=raw.get("embedding"), + metadata=raw.get("metadata") or {}, + ) + + # --------------------------------------------------------------------------- # MongoDBCollection # --------------------------------------------------------------------------- @@ -315,10 +216,9 @@ class MongoDBCollection(AbstractCollection): The ``vectorSearchScore`` returned by Atlas is higher-is-better and is used directly as the vd result ``score`` regardless of metric. - The Atlas vector search index named by ``vector_index`` must exist on - the field ``"embedding"`` of this collection *before* :meth:`search` is - called. This adapter does not create it — see the module docstring for - how to do so via the Atlas UI or API. + The Atlas vector search index named by ``vector_index`` is created + automatically on the first :meth:`search` (see :meth:`_ensure_search_index`); + the adapter then blocks until Atlas reports it queryable. Parameters ---------- @@ -354,12 +254,89 @@ def __init__( self.dimension = dimension self.metric = metric self._vector_index = vector_index + #: Set once the Atlas vector search index is confirmed queryable. + self._index_ready = False @property def native(self) -> PymongoCollection: """The raw ``pymongo.collection.Collection`` handle (escape hatch).""" return self._coll + # ----- Atlas vector search index --------------------------------------- # + + def _ensure_search_index(self) -> None: + """ + Create the Atlas vector search index and wait until it is queryable. + + Lazy and idempotent: invoked on the first :meth:`search`, once the + vector dimension is known. The index is created on the ``"embedding"`` + field with this collection's ``dimension`` and ``metric``; the adapter + then blocks until Atlas reports it queryable. A no-op once ready. + + Raises + ------ + RuntimeError + If the dimension is still unknown, if the deployment is not + Atlas-capable (no ``$listSearchIndexes`` support), or if the index + does not become queryable within the timeout. + """ + if self._index_ready: + return + try: + existing = {idx["name"] for idx in self._coll.list_search_indexes()} + except OperationFailure as exc: + raise RuntimeError( + "This MongoDB deployment does not support Atlas Vector Search. " + "Connect to Atlas or a local 'mongodb-atlas-local' / AtlasCLI " + "deployment — a plain mongod cannot run $vectorSearch." + ) from exc + if self._vector_index not in existing: + if self.dimension is None: + raise RuntimeError( + f"Cannot create the Atlas vector search index for " + f"collection {self.name!r}: the vector dimension is unknown. " + f"Write a document first, or pass dimension= to " + f"create_collection." + ) + from pymongo.operations import SearchIndexModel + + similarity = _METRIC_TO_ATLAS_SIMILARITY.get(self.metric, "cosine") + self._coll.create_search_index( + SearchIndexModel( + definition={ + "fields": [ + { + "type": "vector", + "path": "embedding", + "numDimensions": self.dimension, + "similarity": similarity, + } + ] + }, + name=self._vector_index, + type="vectorSearch", + ) + ) + self._wait_until_queryable() + self._index_ready = True + + def _wait_until_queryable( + self, *, timeout: float = 120.0, poll_interval: float = 1.0 + ) -> None: + """Block until the vector search index reports ``queryable``.""" + import time + + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for idx in self._coll.list_search_indexes(self._vector_index): + if idx.get("queryable"): + return + time.sleep(poll_interval) + raise RuntimeError( + f"Atlas vector search index {self._vector_index!r} for collection " + f"{self.name!r} did not become queryable within {timeout:.0f}s." + ) + # ----- raw primitives -------------------------------------------------- # def _write(self, doc: Document) -> None: @@ -472,10 +449,16 @@ def _query( ``vectorSearchScore``, higher-is-better. """ index = kwargs.get("index", self._vector_index) - mongo_filter = _to_mongo_filter(filter) - pipeline = _build_vector_search_pipeline( - vector, limit=limit, index=index, mongo_filter=mongo_filter - ) + # Auto-create/await our managed index; a caller-supplied index= is + # assumed to be managed externally and is used as-is. + if index == self._vector_index: + self._ensure_search_index() + + # Over-fetch when filtering, then filter client-side — Atlas can only + # pre-filter on index-declared fields, so the canonical vd evaluator is + # applied locally (identical to every other server-backed adapter). + fetch = overfetch_limit(limit, filter) + pipeline = _build_vector_search_pipeline(vector, limit=fetch, index=index) # Allow callers to override numCandidates via kwargs. if "num_candidates" in kwargs: @@ -483,17 +466,16 @@ def _query( kwargs["num_candidates"] ) - results = [] - for raw in self._coll.aggregate(pipeline): - results.append( - { - "id": raw["_id"], - "text": raw.get("text", ""), - "score": raw.get("score", 0.0), - "metadata": raw.get("metadata", {}), - } - ) - return results + results = [ + { + "id": raw["_id"], + "text": raw.get("text", ""), + "score": raw.get("score", 0.0), + "metadata": raw.get("metadata", {}), + } + for raw in self._coll.aggregate(pipeline) + ] + return apply_client_filter(results, filter, limit=limit) # --------------------------------------------------------------------------- @@ -512,10 +494,9 @@ class MongoDBClient(AbstractClient): named by ``database`` (default ``"vd"``). The Atlas vector search index on each collection's ``"embedding"`` field - **must be created out-of-band** (Atlas UI, Admin API, or Terraform) before - :meth:`~MongoDBCollection.search` works. Use the ``vector_index`` name - you gave that index when constructing this client (or override per-query - with ``collection.search(..., index="my_index")``). + is created automatically on the first search, named by ``vector_index`` + (default ``"vector_index"``). Pass ``index="my_index"`` to + ``collection.search(...)`` to use a different, externally-managed index. Parameters ---------- @@ -554,8 +535,9 @@ class MongoDBClient(AbstractClient): client = vd.connect("mongodb") # reads MONGODB_URI from env col = client.get_or_create_collection("my_docs", dimension=1536) - # Create the Atlas vector search index on "embedding" before search! col["doc1"] = vd.Document(id="doc1", text="hello", vector=[0.1]*1536) + # The Atlas vector search index is created on the first search call; + # that call blocks until Atlas reports the index queryable. for hit in col.search([0.1]*1536, limit=5): print(hit["id"], hit["score"]) """ diff --git a/vd/backends/pgvector.py b/vd/backends/pgvector.py index 9e94e2f..1a4c35f 100644 --- a/vd/backends/pgvector.py +++ b/vd/backends/pgvector.py @@ -224,11 +224,13 @@ def _ensure_table(self, dimension: int) -> None: """ if self._created: return + # The vector(N) type modifier must be a SQL literal — Postgres rejects + # a bound parameter there ("type modifiers must be simple constants"). + # int() guards the interpolation; dimension is already an int. self._conn.execute( f"CREATE TABLE IF NOT EXISTS {self._tbl} " f"(doc_id text PRIMARY KEY, text text, metadata jsonb, " - f"embedding vector(%s))", - [dimension], + f"embedding vector({int(dimension)}))" ) try: ops = _HNSW_OPS.get(self.metric, "vector_cosine_ops") @@ -266,15 +268,17 @@ def _write_many(self, docs: list[Document]) -> None: if not docs: return self._ensure_table(self.dimension) - self._conn.executemany( - f"INSERT INTO {self._tbl}(doc_id, text, metadata, embedding) " - f"VALUES (%s, %s, %s, %s) " - f"ON CONFLICT (doc_id) DO UPDATE SET " - f"text = EXCLUDED.text, " - f"metadata = EXCLUDED.metadata, " - f"embedding = EXCLUDED.embedding", - [[d.id, d.text, json.dumps(d.metadata or {}), d.vector] for d in docs], - ) + # executemany is a cursor method in psycopg 3 (not a connection method). + with self._conn.cursor() as cur: + cur.executemany( + f"INSERT INTO {self._tbl}(doc_id, text, metadata, embedding) " + f"VALUES (%s, %s, %s, %s) " + f"ON CONFLICT (doc_id) DO UPDATE SET " + f"text = EXCLUDED.text, " + f"metadata = EXCLUDED.metadata, " + f"embedding = EXCLUDED.embedding", + [[d.id, d.text, json.dumps(d.metadata or {}), d.vector] for d in docs], + ) self._conn.commit() def _read(self, key: str) -> Document: @@ -331,9 +335,13 @@ def _query( return [] op = _DISTANCE_OP.get(self.metric, "<=>") fetch = overfetch_limit(limit, filter) + # Cast the bound parameter to ``vector`` explicitly: psycopg sends a + # Python list as ``double precision[]``, and no ``<=>`` operator exists + # between ``vector`` and an array. (On INSERT the column type drives an + # implicit cast, but a bare operand has no such context.) rows = self._conn.execute( f"SELECT doc_id, text, metadata, " - f"embedding {op} %s AS dist " + f"embedding {op} %s::vector AS dist " f"FROM {self._tbl} " f"ORDER BY dist " f"LIMIT %s", diff --git a/vd/backends/sqlite_vec.py b/vd/backends/sqlite_vec.py index 8493eef..dd2e4bf 100644 --- a/vd/backends/sqlite_vec.py +++ b/vd/backends/sqlite_vec.py @@ -104,6 +104,14 @@ def _ensure_tables(self, dimension: int) -> None: ) self._conn.commit() + def _tables_exist(self) -> bool: + """Return ``True`` once the docs + vec0 tables have been created.""" + row = self._conn.execute( + "SELECT 1 FROM sqlite_master WHERE name = ?", + (f"{self.name}_docs",), + ).fetchone() + return row is not None + # ----- raw primitives ------------------------------------------------- # def _write(self, doc: Document) -> None: @@ -116,8 +124,12 @@ def _write(self, doc: Document) -> None: (doc.id, doc.text, json.dumps(doc.metadata or {})), ) rowid = cur.fetchone()[0] + # vec0 virtual tables do not support INSERT OR REPLACE conflict + # resolution on the rowid primary key, so upsert the vector as an + # explicit delete-then-insert (a no-op DELETE on first write). + self._conn.execute(f"DELETE FROM {self._vec_tbl} WHERE rowid = ?", (rowid,)) self._conn.execute( - f"INSERT OR REPLACE INTO {self._vec_tbl}(rowid, embedding) VALUES(?, ?)", + f"INSERT INTO {self._vec_tbl}(rowid, embedding) VALUES(?, ?)", (rowid, sqlite_vec.serialize_float32(doc.vector)), ) self._conn.commit() @@ -153,19 +165,17 @@ def _drop(self, key: str) -> None: self._conn.commit() def _keys(self) -> Iterator[str]: - try: - rows = self._conn.execute(f"SELECT doc_id FROM {self._docs_tbl}").fetchall() - except sqlite3.OperationalError: + if not self._tables_exist(): return iter(()) # tables not created yet (empty collection) + rows = self._conn.execute(f"SELECT doc_id FROM {self._docs_tbl}").fetchall() return iter(r[0] for r in rows) def _count(self) -> int: - try: - return self._conn.execute( - f"SELECT COUNT(*) FROM {self._docs_tbl}" - ).fetchone()[0] - except sqlite3.OperationalError: - return 0 + if not self._tables_exist(): + return 0 # tables not created yet (empty collection) + return self._conn.execute(f"SELECT COUNT(*) FROM {self._docs_tbl}").fetchone()[ + 0 + ] def _query( self, @@ -175,21 +185,32 @@ def _query( filter: Optional[Filter], **kwargs, ) -> Iterable[SearchResult]: - fetch = overfetch_limit(limit, filter) - try: - rows = self._conn.execute( - f"SELECT d.doc_id, d.text, d.metadata, v.distance " - f"FROM {self._vec_tbl} v JOIN {self._docs_tbl} d ON d.rowid = v.rowid " - f"WHERE v.embedding MATCH ? ORDER BY v.distance LIMIT ?", - (sqlite_vec.serialize_float32(vector), fetch), - ).fetchall() - except sqlite3.OperationalError: + if not self._tables_exist(): return [] # empty collection — tables not created yet + fetch = overfetch_limit(limit, filter) + # The vec0 KNN constraint (``k = ?``) must sit directly on the virtual + # table; a bare ``LIMIT`` on a JOIN is not pushed down and sqlite-vec + # rejects the query. Do the KNN in a subquery, then join for payloads. + rows = self._conn.execute( + f"SELECT d.doc_id, d.text, d.metadata, knn.distance " + f"FROM (SELECT rowid, distance FROM {self._vec_tbl} " + f" WHERE embedding MATCH ? AND k = ?) knn " + f"JOIN {self._docs_tbl} d ON d.rowid = knn.rowid " + f"ORDER BY knn.distance", + (sqlite_vec.serialize_float32(vector), fetch), + ).fetchall() results = [ { "id": doc_id, "text": text or "", - "score": score_from_distance(distance, self.metric), + # sqlite-vec yields a NULL distance for a degenerate (zero-norm) + # vector under cosine; treat that as 0.0 similarity, matching + # the memory backend's zero-vector semantics. + "score": ( + score_from_distance(distance, self.metric) + if distance is not None + else 0.0 + ), "metadata": json.loads(metadata or "{}"), } for doc_id, text, metadata, distance in rows diff --git a/vd/backends/weaviate.py b/vd/backends/weaviate.py index ecf57d7..c4fdb15 100644 --- a/vd/backends/weaviate.py +++ b/vd/backends/weaviate.py @@ -278,8 +278,9 @@ def _wcol(self) -> Any: def _class_exists(self) -> bool: """Return ``True`` if the Weaviate class for this collection exists.""" try: - self._weaviate_client.collections.get(self._class_name) - return True + # collections.get() returns a lazy handle without a server check — + # collections.exists() is the actual existence query. + return bool(self._weaviate_client.collections.exists(self._class_name)) except Exception: return False @@ -472,8 +473,9 @@ def __init__( def _class_exists(self, class_name: str) -> bool: """Return ``True`` if the Weaviate class ``class_name`` currently exists.""" try: - self._client.collections.get(class_name) - return True + # collections.get() returns a lazy handle without a server check — + # collections.exists() is the actual existence query. + return bool(self._client.collections.exists(class_name)) except Exception: return False