vecspan 0.1.0__py3-none-any.whl

vecspan/backends/milvus.py

@@ -0,0 +1,696 @@
+"""Milvus adapter built on `pymilvus.MilvusClient`.
+
+Partial payload and vector updates use Milvus `upsert(partial_update=True)`
+when that argument exists. Otherwise the adapter reads the full point, merges
+the change, and upserts it. That read-modify-write path can lose a concurrent
+update that lands between the read and the write.
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+from typing import Any, Mapping, Sequence
+from urllib.parse import unquote, urlsplit
+
+from vecspan.backends.base import VectorBackend
+from vecspan.errors import (
+    BackendError,
+    CollectionAlreadyExistsError,
+    CollectionNotFoundError,
+    ValidationError,
+    VecspanError,
+)
+from vecspan.filters import Filter, compile_milvus
+from vecspan.naming import check_identifier
+from vecspan.types import (
+    CollectionSpec,
+    Distance,
+    IdType,
+    PayloadType,
+    Point,
+    ScrollPage,
+    SearchHit,
+    as_plain,
+    as_vector,
+    spec_from_dict,
+)
+from vecspan.validate import coerce_stored_id, prepare_search_options
+
+_METRIC = {
+    Distance.COSINE: "COSINE",
+    Distance.DOT: "IP",
+    Distance.L2: "L2",
+}
+_METRIC_TO_DISTANCE = {value: key for key, value in _METRIC.items()}
+_PAYLOAD_TYPE = {
+    PayloadType.KEYWORD: "VARCHAR",
+    PayloadType.INT: "INT64",
+    PayloadType.FLOAT: "DOUBLE",
+    PayloadType.BOOL: "BOOL",
+}
+_PAYLOAD_FROM_MILVUS = {
+    "VARCHAR": PayloadType.KEYWORD,
+    "INT64": PayloadType.INT,
+    "INT32": PayloadType.INT,
+    "INT16": PayloadType.INT,
+    "INT8": PayloadType.INT,
+    "FLOAT": PayloadType.FLOAT,
+    "DOUBLE": PayloadType.FLOAT,
+    "BOOL": PayloadType.BOOL,
+}
+_OPTIONS = {"token", "db_name", "timeout"}
+
+
+class MilvusBackend(VectorBackend):
+    """Milvus collections addressed through the MilvusClient RPC API."""
+
+    name = "milvus"
+
+    def __init__(self, client: Any) -> None:
+        self._client = client
+        self._specs: dict[str, CollectionSpec] = {}
+        self._closed = False
+        self._partial: bool | None = None
+
+    @classmethod
+    def from_url(cls, url: str, **kwargs: Any) -> MilvusBackend:
+        """Open a client from a `milvus://` URL."""
+        params = parse_milvus_target(url, **kwargs)
+        try:
+            from pymilvus import MilvusClient
+        except ImportError as exc:
+            raise BackendError("Milvus support requires the milvus extra: pip install 'vecspan[milvus]'") from exc
+        try:
+            client = MilvusClient(**params)
+        except Exception as exc:
+            raise BackendError(f"could not connect to Milvus: {exc}") from exc
+        return cls(client)
+
+    def create_collection(self, spec: CollectionSpec) -> None:
+        def run() -> None:
+            if self._client.has_collection(spec.name):
+                raise CollectionAlreadyExistsError(f"collection {spec.name!r} already exists")
+            _apply_plan(self._client, collection_plan(spec))
+            self._specs[spec.name] = spec
+
+        self._run(run)
+
+    def drop_collection(self, name: str) -> None:
+        def run() -> None:
+            if not self._client.has_collection(name):
+                raise CollectionNotFoundError(f"collection {name!r} does not exist")
+            self._client.drop_collection(name)
+            self._specs.pop(name, None)
+
+        self._run(run)
+
+    def has_collection(self, name: str) -> bool:
+        return self._run(lambda: bool(self._client.has_collection(name)))
+
+    def describe_collection(self, name: str) -> CollectionSpec:
+        return self._run(lambda: self._describe(name))
+
+    def list_collections(self) -> list[str]:
+        return self._run(lambda: [str(name) for name in self._client.list_collections()])
+
+    def upsert(self, name: str, points: Sequence[Point]) -> None:
+        def run() -> None:
+            self._cached(name)
+            rows = [point_to_row(point) for point in points]
+            if rows:
+                self._client.upsert(collection_name=name, data=rows)
+
+        self._run(run)
+
+    def update_payload(self, name: str, ids: Sequence[int | str], payload: dict) -> None:
+        def run() -> None:
+            spec = self._cached(name)
+            if self._supports_partial_update():
+                self._require_existing(name, spec, ids)
+                rows = [{"id": point_id, **payload} for point_id in ids]
+                self._client.upsert(collection_name=name, data=rows, partial_update=True)
+                return
+            self._read_modify_write(name, spec, ids, payload_patch=payload, vectors_by_id={})
+
+        self._run(run)
+
+    def update_vectors(
+        self,
+        name: str,
+        ids: Sequence[int | str],
+        vectors: Sequence[dict[str, list[float]]],
+    ) -> None:
+        def run() -> None:
+            spec = self._cached(name)
+            pairs = list(zip(ids, vectors))
+            if self._supports_partial_update():
+                self._require_existing(name, spec, ids)
+                rows = [{"id": point_id, **vector} for point_id, vector in pairs]
+                self._client.upsert(collection_name=name, data=rows, partial_update=True)
+                return
+            self._read_modify_write(
+                name,
+                spec,
+                ids,
+                payload_patch={},
+                vectors_by_id={point_id: vector for point_id, vector in pairs},
+            )
+
+        self._run(run)
+
+    def delete(
+        self,
+        name: str,
+        *,
+        ids: Sequence[int | str] | None = None,
+        filter: Filter | None = None,
+    ) -> None:
+        def run() -> None:
+            self._cached(name)
+            if ids is not None:
+                self._client.delete(collection_name=name, ids=list(ids))
+                return
+            if filter is None:
+                raise ValidationError("pass exactly one of ids or filter")
+            self._client.delete(collection_name=name, filter=compile_milvus(filter))
+
+        self._run(run)
+
+    def get(self, name: str, ids: Sequence[int | str], *, with_vectors: bool = False) -> list[Point]:
+        def run() -> list[Point]:
+            return self._load(name, ids, with_vectors=with_vectors, ordered=True)
+
+        return self._run(run)
+
+    def count(self, name: str, *, filter: Filter | None = None) -> int:
+        def run() -> int:
+            spec = self._cached(name)
+            expr = compile_milvus(filter) if filter is not None else match_all_expr(spec)
+            rows = self._client.query(collection_name=name, filter=expr, output_fields=["count(*)"])
+            if not rows:
+                return 0
+            row = rows[0]
+            if isinstance(row, Mapping):
+                for key, value in row.items():
+                    if "count" in str(key):
+                        return int(value)
+            raise BackendError("Milvus count response has no count")
+
+        return self._run(run)
+
+    def scroll(
+        self,
+        name: str,
+        *,
+        filter: Filter | None = None,
+        limit: int = 100,
+        offset: int | str | None = None,
+        with_vectors: bool = False,
+    ) -> ScrollPage:
+        def run() -> ScrollPage:
+            spec = self._cached(name)
+            start = 0 if offset is None else offset
+            if isinstance(start, bool) or not isinstance(start, int) or start < 0:
+                raise ValidationError("Milvus scroll offset must be an integer from a previous page")
+            expr = compile_milvus(filter) if filter is not None else match_all_expr(spec)
+            rows = self._client.query(
+                collection_name=name,
+                filter=expr,
+                limit=limit,
+                offset=start,
+                output_fields=_output_fields(spec, with_vectors),
+            )
+            points = [row_to_point(row, spec, with_vectors=with_vectors) for row in rows]
+            next_offset = None if len(points) < limit else start + len(points)
+            return ScrollPage(points=points, next_offset=next_offset)
+
+        return self._run(run)
+
+    def search(
+        self,
+        name: str,
+        vector: Sequence[float],
+        *,
+        vector_name: str,
+        filter: Filter | None = None,
+        limit: int = 10,
+        with_vectors: bool = False,
+        backend_options: dict | None = None,
+    ) -> list[SearchHit]:
+        def run() -> list[SearchHit]:
+            spec = self._cached(name)
+            field = spec.vector(vector_name)
+            options = prepare_search_options(backend_options, allow_nprobe=True)
+            arguments: dict[str, Any] = {
+                "collection_name": name,
+                "data": [list(vector)],
+                "anns_field": field.name,
+                "search_params": milvus_search_params(field.distance, limit, options),
+                "limit": limit,
+                "output_fields": _output_fields(spec, with_vectors),
+            }
+            if filter is not None:
+                arguments["filter"] = compile_milvus(filter)
+            results = self._client.search(**arguments)
+            return parse_search_hits(results, spec, field.name, with_vectors=with_vectors)
+
+        return self._run(run)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        closer = getattr(self._client, "close", None)
+        if not callable(closer):
+            return
+        try:
+            closer()
+        except Exception as exc:
+            raise BackendError(f"milvus close failed: {exc}") from exc
+
+    def _describe(self, name: str) -> CollectionSpec:
+        if not self._client.has_collection(name):
+            raise CollectionNotFoundError(f"collection {name!r} does not exist")
+        description = dict(self._client.describe_collection(name))
+        description.setdefault("collection_name", name)
+        indexes: dict[str, Mapping[str, Any]] = {}
+        for index_name in self._client.list_indexes(name):
+            info = self._client.describe_index(name, index_name)
+            if not isinstance(info, Mapping):
+                info = dict(info)
+            field_name = str(info.get("field_name") or index_name)
+            indexes[field_name] = info
+        spec = spec_from_milvus(description, indexes)
+        self._specs[name] = spec
+        return spec
+
+    def _cached(self, name: str) -> CollectionSpec:
+        spec = self._specs.get(name)
+        if spec is None:
+            return self._describe(name)
+        return spec
+
+    def _supports_partial_update(self) -> bool:
+        if self._partial is None:
+            try:
+                parameters = inspect.signature(self._client.upsert).parameters
+            except (TypeError, ValueError):
+                self._partial = False
+            else:
+                # pymilvus documents partial_update as a **kwargs option.
+                self._partial = "partial_update" in parameters or any(
+                    parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
+                )
+        return self._partial
+
+    def _require_existing(self, name: str, spec: CollectionSpec, ids: Sequence[int | str]) -> None:
+        found = {point.id for point in self._load(name, ids, with_vectors=False, ordered=False)}
+        for point_id in ids:
+            if point_id not in found:
+                raise ValidationError(f"unknown point id: {point_id!r}")
+
+    def _load(
+        self,
+        name: str,
+        ids: Sequence[int | str],
+        *,
+        with_vectors: bool,
+        ordered: bool,
+    ) -> list[Point]:
+        spec = self._cached(name)
+        rows = self._client.get(
+            collection_name=name,
+            ids=list(ids),
+            output_fields=_output_fields(spec, with_vectors),
+        )
+        by_id: dict[int | str, Point] = {}
+        for row in rows:
+            point = row_to_point(row, spec, with_vectors=with_vectors)
+            by_id[point.id] = point
+        if not ordered:
+            return list(by_id.values())
+        return [by_id[point_id] for point_id in ids if point_id in by_id]
+
+    def _read_modify_write(
+        self,
+        name: str,
+        spec: CollectionSpec,
+        ids: Sequence[int | str],
+        *,
+        payload_patch: Mapping[str, Any],
+        vectors_by_id: Mapping[int | str, Mapping[str, Sequence[float]]],
+    ) -> None:
+        current = {
+            point.id: point
+            for point in self._load(name, ids, with_vectors=True, ordered=False)
+        }
+        rows = []
+        for point_id in ids:
+            point = current.get(point_id)
+            if point is None:
+                raise ValidationError(f"unknown point id: {point_id!r}")
+            payload = dict(point.payload)
+            payload.update(payload_patch)
+            vectors = {key: list(value) for key, value in point.vectors.items()}
+            extra = vectors_by_id.get(point_id)
+            if extra:
+                vectors.update({key: list(value) for key, value in extra.items()})
+            rows.append(point_to_row(Point(id=point_id, vectors=vectors, payload=payload)))
+        self._client.upsert(collection_name=name, data=rows)
+
+    def _run(self, action):
+        if self._closed:
+            raise BackendError("client is closed")
+        try:
+            return action()
+        except VecspanError:
+            raise
+        except Exception as exc:
+            raise BackendError(f"milvus request failed: {exc}") from exc
+
+
+def parse_milvus_target(url: str, **kwargs: Any) -> dict[str, Any]:
+    """Return `MilvusClient` keyword arguments for a `milvus://` URL."""
+    unknown = sorted(set(kwargs) - _OPTIONS)
+    if unknown:
+        raise ValidationError(f"unknown option: {unknown[0]}")
+    parsed = urlsplit(url)
+    if parsed.scheme.lower() != "milvus":
+        raise ValidationError("Milvus URL must start with milvus://")
+    if parsed.query or parsed.fragment:
+        raise ValidationError("milvus URL does not take a query string")
+    host = parsed.hostname
+    if not host:
+        raise ValidationError("milvus URL requires a host")
+    port = parsed.port or 19530
+    token = _milvus_token(parsed, kwargs.get("token"))
+    db_name = _milvus_db(parsed.path, kwargs.get("db_name"))
+    params: dict[str, Any] = {"uri": f"http://{host}:{port}"}
+    if token is not None:
+        params["token"] = token
+    if db_name is not None:
+        params["db_name"] = db_name
+    timeout = kwargs.get("timeout")
+    if timeout is not None:
+        if isinstance(timeout, bool) or not isinstance(timeout, (int, float)) or timeout < 0:
+            raise ValidationError("timeout must be a non-negative number")
+        params["timeout"] = float(timeout)
+    return params
+
+
+def collection_plan(spec: CollectionSpec) -> dict[str, Any]:
+    """Describe the Milvus schema and HNSW indexes for `spec`."""
+    return {
+        "collection_name": spec.name,
+        "primary": {
+            "name": "id",
+            "type": "INT64" if spec.id_type == IdType.INT else "VARCHAR",
+            "max_length": 36 if spec.id_type == IdType.UUID else None,
+        },
+        "vectors": [
+            {
+                "name": item.name,
+                "dim": item.size,
+                "metric": _METRIC[item.distance],
+                "m": item.m,
+                "ef_construct": item.ef_construct,
+            }
+            for item in spec.vectors
+        ],
+        "payload": [
+            {
+                "name": item.name,
+                "type": _PAYLOAD_TYPE[item.type],
+                "max_length": item.max_length if item.type == PayloadType.KEYWORD else None,
+                "indexed": item.indexed,
+            }
+            for item in spec.payload_fields
+        ],
+    }
+
+
+def match_all_expr(spec: CollectionSpec) -> str:
+    """Milvus expression that matches every id this library writes."""
+    if spec.id_type == IdType.INT:
+        return "id >= 0"
+    return 'id != ""'
+
+
+def milvus_search_params(distance: Distance, limit: int, options: Mapping[str, int]) -> dict[str, Any]:
+    """Build the HNSW search params passed to MilvusClient.search."""
+    params: dict[str, int] = {"ef": options.get("ef", max(64, limit))}
+    if "nprobe" in options:
+        params["nprobe"] = options["nprobe"]
+    return {"metric_type": _METRIC[distance], "params": params}
+
+
+def point_to_row(point: Point) -> dict[str, Any]:
+    """Flatten a point into the dict MilvusClient upsert expects."""
+    row: dict[str, Any] = {"id": point.id}
+    row.update(point.vectors)
+    row.update(point.payload)
+    return row
+
+
+def row_to_point(row: Mapping[str, Any] | Any, spec: CollectionSpec, *, with_vectors: bool) -> Point:
+    """Split a Milvus row or search hit into a point."""
+    data = _as_mapping(row)
+    entity = data.pop("entity", None)
+    if isinstance(entity, Mapping):
+        for key, value in entity.items():
+            data.setdefault(key, value)
+    meta = data.pop("$meta", None)
+    if isinstance(meta, Mapping):
+        for key, value in meta.items():
+            data.setdefault(key, value)
+    data = as_plain(data)
+    if "id" not in data:
+        raise BackendError("Milvus row is missing id")
+    point_id = coerce_stored_id(spec, data.pop("id"))
+    data.pop("distance", None)
+    data.pop("score", None)
+    vectors: dict[str, list[float]] = {}
+    for field in spec.vectors:
+        if field.name not in data:
+            continue
+        raw = data.pop(field.name)
+        if with_vectors:
+            vectors[field.name] = as_vector(raw)
+    return Point(id=point_id, vectors=vectors if with_vectors else {}, payload=dict(data))
+
+
+def parse_search_hits(
+    results: object,
+    spec: CollectionSpec,
+    vector_name: str,
+    *,
+    with_vectors: bool,
+) -> list[SearchHit]:
+    """Convert a MilvusClient.search result into search hits."""
+    field = spec.vector(vector_name)
+    if isinstance(results, list) and results and isinstance(results[0], list):
+        hits = results[0]
+    elif isinstance(results, list):
+        hits = results
+    else:
+        raise BackendError("unexpected Milvus search response")
+    parsed: list[SearchHit] = []
+    for hit in hits:
+        data = _as_mapping(hit)
+        score = data.get("distance", data.get("score"))
+        if score is None:
+            raise BackendError("Milvus search hit is missing a score")
+        point = row_to_point(data, spec, with_vectors=with_vectors)
+        parsed.append(
+            SearchHit(
+                id=point.id,
+                score=float(score),
+                distance=field.distance,
+                payload=dict(point.payload),
+                vectors=dict(point.vectors) if with_vectors else None,
+            )
+        )
+    return parsed
+
+
+def spec_from_milvus(description: Mapping[str, Any], indexes: Mapping[str, Mapping[str, Any]]) -> CollectionSpec:
+    """Rebuild a collection spec from Milvus describe and index metadata."""
+    primary: tuple[str, str] | None = None
+    vectors = []
+    payload = []
+    for raw in description.get("fields") or []:
+        if not isinstance(raw, Mapping):
+            raw = dict(raw)
+        name = str(raw.get("name"))
+        type_name = _type_name(raw.get("type", raw.get("dtype")))
+        params = _params(raw.get("params"))
+        if raw.get("is_dynamic") or name.startswith("$") or type_name == "JSON":
+            continue
+        is_primary = bool(raw.get("is_primary")) or name == description.get("primary_field")
+        if is_primary:
+            primary = (name, type_name)
+            continue
+        if type_name in {"FLOAT_VECTOR", "FLOAT16_VECTOR", "BFLOAT16_VECTOR"}:
+            if type_name != "FLOAT_VECTOR":
+                raise BackendError(f"unsupported vector type {type_name} on {name!r}")
+            index = dict(indexes.get(name) or {})
+            index_params = _params(index.get("params"))
+            for key in ("M", "efConstruction", "ef_construct"):
+                if key in index and key not in index_params:
+                    index_params[key] = index[key]
+            metric = str(index.get("metric_type") or index_params.get("metric_type") or "").upper()
+            if metric not in _METRIC_TO_DISTANCE:
+                raise BackendError(f"vector field {name!r} has no supported metric")
+            vectors.append(
+                {
+                    "name": name,
+                    "size": int(params.get("dim")),
+                    "distance": _METRIC_TO_DISTANCE[metric],
+                    "m": int(index_params.get("M", 16)),
+                    "ef_construct": int(index_params.get("efConstruction", index_params.get("ef_construct", 200))),
+                }
+            )
+            continue
+        payload_type = _PAYLOAD_FROM_MILVUS.get(type_name)
+        if payload_type is None:
+            raise BackendError(f"unsupported field type {type_name} on {name!r}")
+        payload.append(
+            {
+                "name": name,
+                "type": payload_type,
+                "indexed": name in indexes,
+                "max_length": int(params.get("max_length", 2048)),
+            }
+        )
+    if primary is None:
+        raise BackendError("Milvus collection has no primary key")
+    id_name, id_type_name = primary
+    if id_name != "id":
+        raise BackendError(f"primary key must be named id, got {id_name!r}")
+    if id_type_name == "INT64":
+        id_type = IdType.INT
+    elif id_type_name == "VARCHAR":
+        id_type = IdType.UUID
+    else:
+        raise BackendError(f"unsupported primary key type {id_type_name}")
+    return spec_from_dict(
+        {
+            "name": description.get("collection_name") or description.get("name"),
+            "id_type": id_type.value,
+            "vectors": vectors,
+            "payload_fields": payload,
+        }
+    )
+
+
+def _apply_plan(client: Any, plan: Mapping[str, Any]) -> None:
+    from pymilvus import DataType
+
+    schema = client.create_schema(auto_id=False, enable_dynamic_field=True)
+    primary = plan["primary"]
+    if primary["type"] == "VARCHAR":
+        schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=primary["max_length"])
+    else:
+        schema.add_field("id", DataType.INT64, is_primary=True)
+    for vector in plan["vectors"]:
+        schema.add_field(vector["name"], DataType.FLOAT_VECTOR, dim=vector["dim"])
+    type_map = {
+        "VARCHAR": DataType.VARCHAR,
+        "INT64": DataType.INT64,
+        "DOUBLE": DataType.DOUBLE,
+        "BOOL": DataType.BOOL,
+    }
+    for payload in plan["payload"]:
+        extra = {"max_length": payload["max_length"]} if payload["max_length"] else {}
+        schema.add_field(payload["name"], type_map[payload["type"]], **extra)
+    index_params = client.prepare_index_params()
+    for vector in plan["vectors"]:
+        index_params.add_index(
+            field_name=vector["name"],
+            index_name=vector["name"],
+            index_type="HNSW",
+            metric_type=vector["metric"],
+            params={"M": vector["m"], "efConstruction": vector["ef_construct"]},
+        )
+    for payload in plan["payload"]:
+        if payload["indexed"]:
+            index_params.add_index(
+                field_name=payload["name"],
+                index_name=payload["name"],
+                index_type="INVERTED",
+            )
+    client.create_collection(
+        collection_name=plan["collection_name"],
+        schema=schema,
+        index_params=index_params,
+    )
+
+
+def _output_fields(spec: CollectionSpec, with_vectors: bool) -> list[str]:
+    fields = ["*"]
+    if with_vectors:
+        fields.extend(field.name for field in spec.vectors)
+    return fields
+
+
+def _milvus_token(parsed: Any, explicit: object) -> str | None:
+    if parsed.username or parsed.password:
+        if explicit not in (None, ""):
+            raise ValidationError("pass the Milvus token either in the URL or as token=")
+        user = unquote(parsed.username or "")
+        password = unquote(parsed.password or "")
+        return f"{user}:{password}"
+    if explicit is None:
+        return None
+    if not isinstance(explicit, str):
+        raise ValidationError("token must be a string")
+    return explicit
+
+
+def _milvus_db(path: str, explicit: object) -> str | None:
+    parts = [unquote(part) for part in path.split("/") if part]
+    if explicit is not None and not isinstance(explicit, str):
+        raise ValidationError("db_name must be a string")
+    if not parts:
+        if explicit is None:
+            return None
+        return check_identifier(explicit, "database name")
+    if len(parts) != 1:
+        raise ValidationError("milvus URL path must be a single database name")
+    if explicit is not None and explicit != parts[0]:
+        raise ValidationError("database is set both in the URL and in db_name")
+    return check_identifier(parts[0], "database name")
+
+
+def _type_name(value: object) -> str:
+    if isinstance(value, str):
+        text = value
+    else:
+        named = getattr(value, "name", None)
+        text = named if isinstance(named, str) else str(value)
+    return text.split(".")[-1].upper()
+
+
+def _params(value: object) -> dict[str, Any]:
+    if isinstance(value, str):
+        loaded = json.loads(value)
+        return dict(loaded) if isinstance(loaded, dict) else {}
+    if isinstance(value, Mapping):
+        return dict(value)
+    return {}
+
+
+def _as_mapping(value: object) -> dict[str, Any]:
+    if isinstance(value, Mapping):
+        return dict(value)
+    entity = getattr(value, "entity", None)
+    data = {
+        "id": getattr(value, "id", None),
+        "distance": getattr(value, "distance", getattr(value, "score", None)),
+    }
+    if isinstance(entity, Mapping):
+        data["entity"] = dict(entity)
+    elif entity is not None:
+        data["entity"] = dict(entity)
+    return data