vecspan 0.1.0__py3-none-any.whl

vecspan/types.py

@@ -0,0 +1,282 @@
+"""Public data types shared by every backend."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Mapping, Sequence
+
+from vecspan.errors import ValidationError
+from vecspan.naming import check_collection_name, check_field_name
+
+INT64_MAX = 2**63 - 1
+
+
+class Distance(str, Enum):
+    """How a vector field compares two vectors.
+
+    Cosine and dot scores are larger when the vectors are closer. L2 scores
+    are smaller when they are closer.
+    """
+
+    COSINE = "cosine"
+    DOT = "dot"
+    L2 = "l2"
+
+
+class IdType(str, Enum):
+    """Primary key type chosen when a collection is created."""
+
+    INT = "int"
+    UUID = "uuid"
+
+
+class PayloadType(str, Enum):
+    """Declared scalar type of a payload field."""
+
+    KEYWORD = "keyword"
+    INT = "int"
+    FLOAT = "float"
+    BOOL = "bool"
+
+
+def _as_enum(enum_cls: type[Enum], value: object, label: str) -> Enum:
+    if isinstance(value, enum_cls):
+        return value
+    try:
+        return enum_cls(value)
+    except ValueError:
+        allowed = ", ".join(item.value for item in enum_cls)
+        raise ValidationError(f"invalid {label}: {value!r} (expected one of {allowed})") from None
+
+
+def _as_positive_int(value: object, label: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValidationError(f"{label} must be a positive integer")
+    return value
+
+
+@dataclass(frozen=True)
+class VectorField:
+    """A named dense vector stored on every point."""
+
+    name: str
+    size: int
+    distance: Distance = Distance.COSINE
+    m: int = 16
+    ef_construct: int = 200
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "name", check_field_name(self.name))
+        object.__setattr__(self, "size", _as_positive_int(self.size, "vector size"))
+        object.__setattr__(self, "distance", _as_enum(Distance, self.distance, "distance"))
+        object.__setattr__(self, "m", _as_positive_int(self.m, "m"))
+        if self.m < 2:
+            raise ValidationError("m must be at least 2")
+        object.__setattr__(self, "ef_construct", _as_positive_int(self.ef_construct, "ef_construct"))
+
+
+@dataclass(frozen=True)
+class PayloadField:
+    """A scalar payload field. Indexed fields can be filtered efficiently."""
+
+    name: str
+    type: PayloadType
+    indexed: bool = True
+    max_length: int = 2048
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "name", check_field_name(self.name))
+        object.__setattr__(self, "type", _as_enum(PayloadType, self.type, "payload type"))
+        if not isinstance(self.indexed, bool):
+            raise ValidationError("indexed must be a bool")
+        object.__setattr__(self, "max_length", _as_positive_int(self.max_length, "max_length"))
+        if self.max_length > 65535:
+            raise ValidationError("max_length must be at most 65535")
+
+
+@dataclass(frozen=True)
+class CollectionSpec:
+    """Schema declared when a collection is created and then left unchanged."""
+
+    name: str
+    id_type: IdType
+    vectors: Sequence[VectorField]
+    payload_fields: Sequence[PayloadField] = ()
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "name", check_collection_name(self.name))
+        object.__setattr__(self, "id_type", _as_enum(IdType, self.id_type, "id_type"))
+        object.__setattr__(self, "vectors", _as_fields(self.vectors, VectorField, "vectors"))
+        object.__setattr__(
+            self,
+            "payload_fields",
+            _as_fields(self.payload_fields, PayloadField, "payload_fields"),
+        )
+        if not self.vectors:
+            raise ValidationError("a collection needs at least one vector field")
+        names = [field.name for field in self.vectors]
+        names.extend(field.name for field in self.payload_fields)
+        if len(names) != len(set(names)):
+            raise ValidationError("field names must be unique across vectors and payload")
+
+    def vector(self, name: str | None) -> VectorField:
+        """Return the named vector field, or the only one when `name` is omitted."""
+        if name is None:
+            if len(self.vectors) != 1:
+                raise ValidationError("vector_name is required when the collection has several vectors")
+            return self.vectors[0]
+        for item in self.vectors:
+            if item.name == name:
+                return item
+        raise ValidationError(f"unknown vector field: {name!r}")
+
+    def payload(self, name: str) -> PayloadField | None:
+        """Return the declared payload field, or None when the key is dynamic."""
+        for item in self.payload_fields:
+            if item.name == name:
+                return item
+        return None
+
+
+def _as_fields(value: object, kind: type, label: str) -> tuple[Any, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, kind):
+        raise ValidationError(f"{label} must be a sequence")
+    if isinstance(value, (str, bytes)) or not isinstance(value, Sequence):
+        raise ValidationError(f"{label} must be a sequence")
+    items = tuple(value)
+    for item in items:
+        if not isinstance(item, kind):
+            raise ValidationError(f"{label} must contain {kind.__name__} values")
+    return items
+
+
+@dataclass
+class Point:
+    """One record: an id, one or more named vectors, and a payload mapping."""
+
+    id: int | str
+    vectors: Mapping[str, Sequence[float]]
+    payload: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.vectors, Mapping):
+            raise ValidationError("vectors must be a mapping of field name to vector")
+        if self.payload is None:
+            self.payload = {}
+        elif not isinstance(self.payload, Mapping):
+            raise ValidationError("payload must be a mapping")
+        else:
+            self.payload = dict(self.payload)
+
+
+@dataclass(frozen=True)
+class SearchHit:
+    """One neighbor returned by `search`.
+
+    `score` is the backend's native value. It is not rescaled across backends.
+    Memory and Qdrant report plain Euclidean distance for L2. Milvus reports
+    squared Euclidean distance for L2.
+    """
+
+    id: int | str
+    score: float
+    distance: Distance
+    payload: dict[str, Any]
+    vectors: dict[str, list[float]] | None = None
+
+
+@dataclass(frozen=True)
+class ScrollPage:
+    """A page of points. Pass `next_offset` back to `scroll` to continue.
+
+    The offset is opaque and only meaningful for the same collection on the
+    same backend.
+    """
+
+    points: list[Point]
+    next_offset: int | str | None
+
+
+def spec_to_dict(spec: CollectionSpec) -> dict[str, Any]:
+    """Serialize a spec so a backend can store it next to the collection."""
+    return {
+        "name": spec.name,
+        "id_type": spec.id_type.value,
+        "vectors": [
+            {
+                "name": item.name,
+                "size": item.size,
+                "distance": item.distance.value,
+                "m": item.m,
+                "ef_construct": item.ef_construct,
+            }
+            for item in spec.vectors
+        ],
+        "payload_fields": [
+            {
+                "name": item.name,
+                "type": item.type.value,
+                "indexed": item.indexed,
+                "max_length": item.max_length,
+            }
+            for item in spec.payload_fields
+        ],
+    }
+
+
+def spec_from_dict(data: Mapping[str, Any]) -> CollectionSpec:
+    """Rebuild a spec previously produced by `spec_to_dict`."""
+    try:
+        vectors = [
+            VectorField(
+                name=item["name"],
+                size=item["size"],
+                distance=item["distance"],
+                m=item["m"],
+                ef_construct=item["ef_construct"],
+            )
+            for item in data["vectors"]
+        ]
+        payload_fields = [
+            PayloadField(
+                name=item["name"],
+                type=item["type"],
+                indexed=item["indexed"],
+                max_length=item["max_length"],
+            )
+            for item in data.get("payload_fields", ())
+        ]
+        return CollectionSpec(
+            name=data["name"],
+            id_type=data["id_type"],
+            vectors=vectors,
+            payload_fields=payload_fields,
+        )
+    except (KeyError, TypeError) as exc:
+        raise ValidationError(f"stored collection spec is invalid: {exc}") from exc
+
+
+def as_plain(value: Any) -> Any:
+    """Convert numpy scalars and arrays into plain Python values."""
+    if isinstance(value, dict):
+        return {str(key): as_plain(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [as_plain(item) for item in value]
+    tolist = getattr(value, "tolist", None)
+    if callable(tolist) and not isinstance(value, (str, bytes)):
+        return as_plain(tolist())
+    return value
+
+
+def as_vector(value: Any) -> list[float]:
+    """Return a vector as a list of floats."""
+    data = as_plain(value)
+    if isinstance(data, (str, bytes)) or not isinstance(data, Sequence):
+        raise ValidationError("vector must be a sequence of numbers")
+    try:
+        return [float(item) for item in data]
+    except (TypeError, ValueError) as exc:
+        raise ValidationError("vector must be a sequence of numbers") from exc

vecspan/validate.py

@@ -0,0 +1,348 @@
+"""Turn caller input into the values backends store and query."""
+
+from __future__ import annotations
+
+import math
+import uuid
+from typing import Any, Mapping, Sequence
+
+from vecspan.errors import ValidationError
+from vecspan.filters import And, Cmp, Eq, Expr, Filter, In, IsNull, Not, Or
+from vecspan.naming import check_field_name
+from vecspan.types import (
+    INT64_MAX,
+    CollectionSpec,
+    IdType,
+    PayloadType,
+    Point,
+    VectorField,
+)
+
+PointId = int | str
+
+
+def prepare_points(spec: CollectionSpec, points: object) -> list[Point]:
+    """Validate and normalize a batch of points."""
+    items = _point_list(points)
+    prepared: list[Point] = []
+    seen: set[PointId] = set()
+    for item in items:
+        point = _normalize_point(spec, _as_point(item))
+        if point.id in seen:
+            raise ValidationError(f"duplicate point id: {point.id!r}")
+        seen.add(point.id)
+        prepared.append(point)
+    return prepared
+
+
+def prepare_ids(spec: CollectionSpec, ids: object) -> list[PointId]:
+    """Validate a non-empty id sequence and preserve its order."""
+    items = _id_list(ids)
+    if not items:
+        raise ValidationError("ids must not be empty")
+    prepared = [normalize_id(spec, item) for item in items]
+    if len(prepared) != len(set(prepared)):
+        raise ValidationError("duplicate ids")
+    return prepared
+
+
+def prepare_payload(spec: CollectionSpec, payload: object, *, allow_empty: bool) -> dict[str, Any]:
+    """Validate a payload mapping against declared fields."""
+    if not isinstance(payload, Mapping):
+        raise ValidationError("payload must be a mapping")
+    if not payload and not allow_empty:
+        raise ValidationError("payload must not be empty")
+    vector_names = {item.name for item in spec.vectors}
+    cleaned: dict[str, Any] = {}
+    for key, value in payload.items():
+        name = check_field_name(key)
+        if name in vector_names:
+            raise ValidationError(f"payload key {name!r} collides with a vector field")
+        cleaned[name] = _payload_value(spec, name, value)
+    return cleaned
+
+
+def prepare_vector_updates(
+    spec: CollectionSpec,
+    ids: object,
+    vectors: object,
+) -> list[tuple[PointId, dict[str, list[float]]]]:
+    """Validate parallel ids and partial vector replacements."""
+    prepared_ids = prepare_ids(spec, ids)
+    if isinstance(vectors, Mapping) or not isinstance(vectors, Sequence):
+        raise ValidationError("vectors must be a sequence aligned with ids")
+    if len(vectors) != len(prepared_ids):
+        raise ValidationError("vectors must have one entry per id")
+    updates: list[tuple[PointId, dict[str, list[float]]]] = []
+    for point_id, raw in zip(prepared_ids, vectors):
+        if not isinstance(raw, Mapping) or isinstance(raw, (str, bytes)):
+            raise ValidationError("each vector update must be a mapping of field name to vector")
+        if not raw:
+            raise ValidationError("a vector update must include at least one field")
+        cleaned: dict[str, list[float]] = {}
+        for name, value in raw.items():
+            if not isinstance(name, str):
+                raise ValidationError(f"unknown vector field: {name!r}")
+            field = spec.vector(name)
+            cleaned[field.name] = _vector_values(value, field)
+        updates.append((point_id, cleaned))
+    return updates
+
+
+def prepare_query_vector(
+    spec: CollectionSpec,
+    vector: object,
+    vector_name: str | None,
+) -> tuple[VectorField, list[float]]:
+    """Validate the query vector and resolve which field it belongs to."""
+    field = spec.vector(vector_name)
+    return field, _vector_values(vector, field)
+
+
+def prepare_filter(spec: CollectionSpec, filt: Filter | None) -> Filter | None:
+    """Check filter fields against the collection schema."""
+    if filt is None:
+        return None
+    if not isinstance(filt, Filter):
+        raise ValidationError("filter must be a Filter")
+    _check_expr(spec, filt.expr)
+    return filt
+
+
+def prepare_search_options(options: object, *, allow_nprobe: bool) -> dict[str, int]:
+    """Normalize search tuning keys understood by the backends."""
+    if options is None:
+        return {}
+    if not isinstance(options, Mapping):
+        raise ValidationError("backend_options must be a mapping")
+    allowed = {"ef", "hnsw_ef"}
+    if allow_nprobe:
+        allowed.add("nprobe")
+    unknown = sorted(set(options) - allowed)
+    if unknown:
+        raise ValidationError(f"unknown backend option: {unknown[0]}")
+    if "ef" in options and "hnsw_ef" in options:
+        raise ValidationError("pass only one of ef and hnsw_ef")
+    cleaned: dict[str, int] = {}
+    if "ef" in options or "hnsw_ef" in options:
+        raw = options["ef"] if "ef" in options else options["hnsw_ef"]
+        cleaned["ef"] = _option_int(raw, "ef")
+    if "nprobe" in options:
+        cleaned["nprobe"] = _option_int(options["nprobe"], "nprobe")
+    return cleaned
+
+
+def check_limit(limit: object) -> int:
+    if isinstance(limit, bool) or not isinstance(limit, int) or limit < 1:
+        raise ValidationError("limit must be a positive integer")
+    return limit
+
+
+def check_with_vectors(value: object) -> bool:
+    if not isinstance(value, bool):
+        raise ValidationError("with_vectors must be a bool")
+    return value
+
+
+def normalize_id(spec: CollectionSpec, value: object) -> PointId:
+    """Return the canonical id for this collection."""
+    if isinstance(value, bool) or value is None:
+        raise ValidationError(f"invalid point id: {value!r}")
+    if spec.id_type == IdType.INT:
+        if isinstance(value, int):
+            if value < 0 or value > INT64_MAX:
+                raise ValidationError(f"int id out of range: {value}")
+            return value
+        raise ValidationError(f"expected an int id, got {value!r}")
+    if isinstance(value, uuid.UUID):
+        return str(value)
+    if isinstance(value, str):
+        try:
+            return str(uuid.UUID(value))
+        except ValueError:
+            raise ValidationError(f"expected a UUID id, got {value!r}") from None
+    raise ValidationError(f"expected a UUID id, got {value!r}")
+
+
+def coerce_stored_id(spec: CollectionSpec, value: object) -> PointId:
+    """Normalize an id coming back from a database."""
+    plain = value.item() if callable(getattr(value, "item", None)) and not isinstance(value, (str, bytes)) else value
+    try:
+        return normalize_id(spec, plain)
+    except ValidationError:
+        if spec.id_type == IdType.INT:
+            return int(plain)  # type: ignore[arg-type]
+        return str(plain)
+
+
+def _point_list(points: object) -> list[object]:
+    if isinstance(points, Point) or _is_point_mapping(points):
+        return [points]
+    if isinstance(points, (str, bytes)) or not isinstance(points, Sequence):
+        raise ValidationError("points must be a Point or a sequence of points")
+    return list(points)
+
+
+def _id_list(ids: object) -> list[object]:
+    if isinstance(ids, (str, int)) and not isinstance(ids, bool):
+        return [ids]
+    if isinstance(ids, (str, bytes)) or not isinstance(ids, Sequence):
+        raise ValidationError("ids must be a sequence")
+    return list(ids)
+
+
+def _is_point_mapping(value: object) -> bool:
+    return isinstance(value, Mapping) and "id" in value and "vectors" in value
+
+
+def _as_point(value: object) -> Point:
+    if isinstance(value, Point):
+        return value
+    if isinstance(value, Mapping):
+        if "id" not in value or "vectors" not in value:
+            raise ValidationError("point mapping needs id and vectors")
+        payload = value.get("payload") or {}
+        return Point(id=value["id"], vectors=value["vectors"], payload=payload)
+    raise ValidationError("point must be a Point or a mapping")
+
+
+def _normalize_point(spec: CollectionSpec, point: Point) -> Point:
+    point_id = normalize_id(spec, point.id)
+    if not point.vectors:
+        raise ValidationError("a point needs every vector field")
+    vectors: dict[str, list[float]] = {}
+    for field in spec.vectors:
+        if field.name not in point.vectors:
+            raise ValidationError(f"point is missing vector {field.name!r}")
+        vectors[field.name] = _vector_values(point.vectors[field.name], field)
+    extra = set(point.vectors) - set(vectors)
+    if extra:
+        name = sorted(extra)[0]
+        raise ValidationError(f"unknown vector field: {name!r}")
+    payload = prepare_payload(spec, point.payload, allow_empty=True)
+    return Point(id=point_id, vectors=vectors, payload=payload)
+
+
+def _vector_values(value: object, field: VectorField) -> list[float]:
+    if isinstance(value, (str, bytes)) or not isinstance(value, Sequence):
+        raise ValidationError(f"vector {field.name!r} must be a sequence of numbers")
+    if len(value) != field.size:
+        raise ValidationError(
+            f"vector {field.name!r} has dimension {len(value)}, expected {field.size}"
+        )
+    numbers: list[float] = []
+    for item in value:
+        if isinstance(item, bool) or not isinstance(item, (int, float)):
+            raise ValidationError(f"vector {field.name!r} must contain numbers")
+        number = float(item)
+        if not math.isfinite(number):
+            raise ValidationError(f"vector {field.name!r} must be finite")
+        numbers.append(number)
+    return numbers
+
+
+def _payload_value(spec: CollectionSpec, name: str, value: object) -> Any:
+    declared = spec.payload(name)
+    if declared is None:
+        return _dynamic_value(name, value)
+    if value is None:
+        return None
+    if declared.type == PayloadType.KEYWORD:
+        if not isinstance(value, str):
+            raise ValidationError(f"payload field {name!r} expects a string")
+        if len(value) > declared.max_length:
+            raise ValidationError(f"payload field {name!r} exceeds max_length {declared.max_length}")
+        return value
+    if declared.type == PayloadType.INT:
+        if isinstance(value, bool) or not isinstance(value, int):
+            raise ValidationError(f"payload field {name!r} expects an int")
+        if value < -INT64_MAX - 1 or value > INT64_MAX:
+            raise ValidationError(f"payload field {name!r} is outside int64")
+        return value
+    if declared.type == PayloadType.FLOAT:
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            raise ValidationError(f"payload field {name!r} expects a number")
+        number = float(value)
+        if not math.isfinite(number):
+            raise ValidationError(f"payload field {name!r} must be finite")
+        return number
+    if isinstance(value, bool):
+        return value
+    raise ValidationError(f"payload field {name!r} expects a bool")
+
+
+def _dynamic_value(name: str, value: object) -> Any:
+    if value is None or isinstance(value, (str, bool)):
+        return value
+    if isinstance(value, float):
+        if not math.isfinite(value):
+            raise ValidationError(f"payload field {name!r} must be finite")
+        return value
+    if isinstance(value, int):
+        return value
+    if isinstance(value, list):
+        return [_dynamic_value(name, item) for item in value]
+    if isinstance(value, dict):
+        return {check_field_name(key): _dynamic_value(str(key), item) for key, item in value.items()}
+    raise ValidationError(f"payload field {name!r} has unsupported type {type(value).__name__}")
+
+
+def _check_expr(spec: CollectionSpec, expr: Expr) -> None:
+    if isinstance(expr, (Eq, In, Cmp, IsNull)):
+        _check_field_expr(spec, expr)
+        return
+    if isinstance(expr, Not):
+        _check_expr(spec, expr.inner)
+        return
+    if isinstance(expr, (And, Or)):
+        for child in expr.children:
+            _check_expr(spec, child)
+        return
+    raise ValidationError(f"unsupported filter node: {type(expr).__name__}")
+
+
+def _check_field_expr(spec: CollectionSpec, expr: Eq | In | Cmp | IsNull) -> None:
+    if any(field.name == expr.field for field in spec.vectors):
+        raise ValidationError(f"cannot filter on vector field {expr.field!r}")
+    declared = spec.payload(expr.field)
+    if declared is None or isinstance(expr, IsNull):
+        return
+    if declared.type == PayloadType.KEYWORD:
+        if isinstance(expr, Cmp):
+            raise ValidationError(f"cannot compare string field {expr.field!r}")
+        _expect_kind(expr, "string")
+        return
+    if declared.type == PayloadType.BOOL:
+        if isinstance(expr, Cmp):
+            raise ValidationError(f"cannot compare bool field {expr.field!r}")
+        _expect_kind(expr, "bool")
+        return
+    if declared.type == PayloadType.INT:
+        _expect_kind(expr, "int")
+        return
+    _expect_kind(expr, "number")
+
+
+def _expect_kind(expr: Eq | In | Cmp, kind: str) -> None:
+    values = expr.values if isinstance(expr, In) else (expr.value,)
+    for value in values:
+        if not _matches_kind(value, kind):
+            raise ValidationError(f"filter value for {expr.field!r} does not match the field type")
+
+
+def _matches_kind(value: object, kind: str) -> bool:
+    if kind == "string":
+        return isinstance(value, str)
+    if kind == "bool":
+        return isinstance(value, bool)
+    if kind == "int":
+        return isinstance(value, int) and not isinstance(value, bool)
+    if kind == "number":
+        return isinstance(value, (int, float)) and not isinstance(value, bool)
+    return False
+
+
+def _option_int(value: object, label: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValidationError(f"{label} must be a positive integer")
+    return value