diff --git a/dev/agent-skills/target-connector/SKILL.md b/dev/agent-skills/target-connector/SKILL.md index 1d8c43ae4..07c2c1185 100644 --- a/dev/agent-skills/target-connector/SKILL.md +++ b/dev/agent-skills/target-connector/SKILL.md @@ -217,6 +217,37 @@ _shared_sink = coco.TargetActionSink.from_fn(_apply_actions) When building queries from user-provided names (table, column, index) or values (record IDs, keys), you must guard against injection and ensure correctness. See [input_safety.md](input_safety.md) for patterns on identifier validation, parameterized queries, and value escaping. +### Vector Schema Resolution + +Table-like connectors that infer dense or sparse vector columns from annotations must use the shared connector-kit resolver: + +```python +from cocoindex.connectorkits import resolve_vector_schemas + +schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, + # Include this only when the connector supports dense but not sparse vectors. + reject_sparse_vectors_for="MyConnector", +) +``` + +The resolver gives direct annotations precedence over `ContextKey` providers, detects dense/sparse conflicts, and validates that `SparseVectorSchema` is attached only to a `SparseVector` field. Do not duplicate these checks in the connector. + +Connectors that do not consume vector metadata but still need to reject sparse columns should use the non-resolving guard, which deliberately ignores `ContextKey` annotations: + +```python +from cocoindex.connectorkits import reject_sparse_vectors + +reject_sparse_vectors( + type_info.base_type, + all_annotations, + connector_name="MyConnector", +) +``` + +Do not add `SparseVectorSchemaProvider` to a public override union when the connector rejects sparse vectors; annotations remain runtime-guarded without advertising unsupported input. + ## Completion Checklist After implementing the connector code, complete these additional steps: diff --git a/docs/src/content/docs/common_resources/vector_schema.mdx b/docs/src/content/docs/common_resources/vector_schema.mdx index 9b6481491..f53c95b8d 100644 --- a/docs/src/content/docs/common_resources/vector_schema.mdx +++ b/docs/src/content/docs/common_resources/vector_schema.mdx @@ -1,9 +1,8 @@ --- title: "*Vector schema* annotations" description: > - Describe vector columns with VectorSchema and VectorSchemaProvider. Covers - the three annotation patterns (ContextKey, embedder instance, explicit - schema) and MultiVectorSchema for models like ColBERT. + Describe dense, sparse, and multi-vector columns with shared schema types and + provider protocols. --- The schema module (`cocoindex.resources.schema`) defines types that describe vector columns. CocoIndex connectors use these to automatically configure the correct column type (e.g., `vector(384)` in Postgres, `fixed_size_list(384)` in LanceDB). @@ -119,3 +118,39 @@ multi_schema = MultiVectorSchema( vector_schema=VectorSchema(dtype=np.dtype(np.float32), size=128) ) ``` + +## SparseVector / SparseVectorSchema + +`SparseVector` is the canonical value shared by connectors that support sparse vectors. It stores parallel tuples of 0-based integer indices and finite float weights. Indices must be non-negative, sorted ascending, and unique. Direct construction accepts only that canonical tuple/int/float representation. Use `from_arrays()` to coerce parallel numeric iterables and sort them by index; `from_mapping()` delegates to the same normalization path. `as_sparse_vector()` accepts either an already-canonical value or a mapping. + +```python +from cocoindex.resources.schema import SparseVector + +sparse = SparseVector.from_mapping({7: 0.9, 1: 0.5}) +assert sparse.indices == (1, 7) +assert sparse.values == (0.5, 0.9) + +sparse = SparseVector.from_arrays(encoder.indices, encoder.values) +assert sparse.indices == tuple(sorted(encoder.indices)) +``` + +`SparseVectorSchema` carries optional dimensionality. Sparse values use float32 across CocoIndex connectors. When supplied, `size` must be positive. The dimension is required by Postgres `sparsevec(n)` but optional for backends such as zvec and Qdrant. + +```python +from dataclasses import dataclass +from typing import Annotated + +from cocoindex.resources.schema import SparseVector, SparseVectorSchema + +@dataclass +class SparseDoc: + id: str + sparse: Annotated[ + SparseVector, + SparseVectorSchema(size=50_000), + ] +``` + +`SparseVectorSchemaProvider` is the corresponding provider protocol. Its `__coco_sparse_vector_schema__()` method can be supplied directly or through a `ContextKey`, following the same resolution pattern as `VectorSchemaProvider`. + +Metrics, scoring modifiers, index methods, and quantization are connector-specific and therefore stay on connector configuration rather than `SparseVectorSchema`. diff --git a/docs/src/content/docs/connectors/postgres.mdx b/docs/src/content/docs/connectors/postgres.mdx index a1beb5f13..ac842b35f 100644 --- a/docs/src/content/docs/connectors/postgres.mdx +++ b/docs/src/content/docs/connectors/postgres.mdx @@ -241,6 +241,8 @@ The actual PostgreSQL index is named `{table_name}__vector__{name}`. - `m` — Maximum number of connections per layer (hnsw only). - `ef_construction` — Size of the dynamic candidate list for construction (hnsw only). +`sparsevec` columns support only `method="hnsw"`; passing the default `"ivfflat"` raises an error. pgvector limits an HNSW-indexed sparse vector to 1,000 non-zero elements. + **Example:** ```python @@ -295,7 +297,9 @@ async def TableSchema.from_class( record_type: type[RowT], primary_key: list[str], *, - column_overrides: dict[str, PgType | VectorSchemaProvider] | None = None, + column_overrides: dict[ + str, PgType | VectorSchemaProvider | SparseVectorSchemaProvider + ] | None = None, ) -> TableSchema[RowT] ``` @@ -338,6 +342,7 @@ Python types are automatically mapped to PostgreSQL types: | `datetime.timedelta` | `interval` | | `list`, `dict`, nested structs | `jsonb` | | `NDArray` (with vector schema) | `vector(n)` or `halfvec(n)` | +| `SparseVector` (with sparse vector schema) | `sparsevec(n)` | :::note[U+0000 (NUL) in strings] U+0000 (NUL) is a valid Unicode codepoint, but Postgres cannot store it — neither in `text`-family columns nor inside strings in `jsonb` (the `\u0000` escape is rejected at parse time). CocoIndex automatically strips U+0000 from strings before writing to Postgres, recursively for nested strings and dict keys in `jsonb` payloads. For example, `"Hello\0World"` is written as `"HelloWorld"`. @@ -379,6 +384,31 @@ schema = postgres.TableSchema( For `NDArray` fields, a [`VectorSchemaProvider`](../common_resources/vector_schema#vectorschemaprovider) annotation specifies the vector dimension and dtype. The connector has built-in pgvector support and automatically creates the extension when needed. See [Vector Schema](../common_resources/vector_schema#vectorschemaprovider) for the full list of annotation options (`ContextKey`, embedder instance, or explicit `VectorSchema`). +#### Sparse vectors + +Postgres requires a dimension for `sparsevec(n)`, and its sparse values are float32. Annotate the shared `SparseVector` type with `SparseVectorSchema(size=...)`; CocoIndex converts its canonical 0-based indices to pgvector's 1-based wire format. + +```python +from dataclasses import dataclass +from typing import Annotated + +from cocoindex.resources.schema import SparseVector, SparseVectorSchema + +@dataclass +class SparseDoc: + id: str + sparse: Annotated[ + SparseVector, + SparseVectorSchema(size=50_000), + ] + +schema = await postgres.TableSchema.from_class(SparseDoc, primary_key=["id"]) +table = await postgres.mount_table_target(PG_DB, "sparse_docs", schema) +table.declare_vector_index(column="sparse", metric="ip", method="hnsw") +``` + +Rows may provide a `SparseVector` or an integer-to-float mapping at declaration time; mappings are validated and sorted before encoding. A bare `SparseVector` field or a schema without `size` raises an actionable error. + ### Table schema: explicit column definitions Define columns directly using `ColumnDef`: diff --git a/docs/src/content/docs/connectors/qdrant.mdx b/docs/src/content/docs/connectors/qdrant.mdx index 9078978ed..06a3f21a2 100644 --- a/docs/src/content/docs/connectors/qdrant.mdx +++ b/docs/src/content/docs/connectors/qdrant.mdx @@ -231,7 +231,6 @@ point = qdrant.PointStruct( For hybrid retrieval, declare a named dense vector and a named sparse vector in the same collection: ```python -from qdrant_client.http import models as qdrant_models from cocoindex.resources.schema import VectorSchema import numpy as np @@ -246,23 +245,22 @@ schema = await qdrant.CollectionSchema.create( ) ``` -Points put both vector types in `PointStruct.vector`. The sparse vector is a native Qdrant `SparseVector`, not payload: +Points put both vector types in `PointStruct.vector`. Use `qdrant.sparse_vector()` to convert the shared canonical `SparseVector` (or an integer-to-float mapping) to Qdrant's native model: ```python point = qdrant.PointStruct( - id=123, + id=doc_id, vector={ "dense": dense_embedding.tolist(), - "sparse": qdrant_models.SparseVector( - indices=sparse_indices, - values=sparse_values, - ), + "sparse": qdrant.sparse_vector({7: 0.9, 1: 0.5}), }, payload={"text": text}, ) target.declare_point(point) ``` +Qdrant does not require sparse dimensionality, so `SparseVectorSchema` is not needed here; a shared [`SparseVector`](../common_resources/vector_schema#sparsevector--sparsevectorschema) annotation on your row type remains useful when the same rows also target Postgres or zvec. Sparse scoring configuration, including `modifier="idf"`, stays on `QdrantSparseVectorDef`. + #### VectorSchemaProvider The `schema` field of `QdrantVectorDef` accepts a [`VectorSchemaProvider`](../common_resources/vector_schema#vectorschemaprovider), a `ContextKey`, or an explicit `VectorSchema` to specify the vector dimension and dtype. See [Vector Schema](../common_resources/vector_schema#vectorschemaprovider) for details. @@ -391,8 +389,7 @@ async def app_main() -> None: ### Example: dense + sparse hybrid vectors ```python -from qdrant_client.http import models as qdrant_models -from cocoindex.resources.schema import VectorSchema +from cocoindex.resources.schema import SparseVector, VectorSchema import numpy as np @coco.fn @@ -416,9 +413,11 @@ async def app_main() -> None: id=chunk.id, vector={ "dense": chunk.dense_embedding.tolist(), - "sparse": qdrant_models.SparseVector( - indices=chunk.sparse_indices, - values=chunk.sparse_values, + "sparse": qdrant.sparse_vector( + SparseVector.from_arrays( + chunk.sparse_indices, + chunk.sparse_values, + ) ), }, payload={"text": chunk.text}, @@ -508,9 +507,8 @@ results = client.query_points( limit=100, ), qdrant_models.Prefetch( - query=qdrant_models.SparseVector( - indices=query_sparse_indices, - values=query_sparse_values, + query=qdrant.sparse_vector( + dict(zip(query_sparse_indices, query_sparse_values, strict=True)) ), using="sparse", limit=100, diff --git a/docs/src/content/docs/connectors/zvec.mdx b/docs/src/content/docs/connectors/zvec.mdx index 00e4a7638..96a7fb92b 100644 --- a/docs/src/content/docs/connectors/zvec.mdx +++ b/docs/src/content/docs/connectors/zvec.mdx @@ -122,7 +122,11 @@ async def CollectionSchema.from_class( record_type: type[RowT], primary_key: list[str], *, - column_overrides: dict[str, ZvecType | ZvecVectorDef | ZvecFtsType | VectorSchemaProvider] | None = None, + column_overrides: dict[ + str, + ZvecType | ZvecVectorDef | ZvecFtsType + | VectorSchemaProvider | SparseVectorSchemaProvider, + ] | None = None, ) -> CollectionSchema[RowT] ``` @@ -175,6 +179,7 @@ Scalar Python types map to zvec field types as follows: | `list[str]` / `list[int]` / `list[float]` / `list[bool]` | `ARRAY_STRING` / `ARRAY_INT64` / `ARRAY_DOUBLE` / `ARRAY_BOOL` | | other `list`, `dict`, nested structs | `STRING` (JSON) | | `NDArray` (with vector schema) | `VECTOR_FP32` (float32) or `VECTOR_FP16` (float16) | +| `SparseVector` | `SPARSE_VECTOR_FP32` | Scalar fields get an invert index by default so they can be used in query filters. The primary-key column maps to the document `id` and is not stored as a separate field. @@ -237,15 +242,27 @@ class Doc: #### Sparse vectors -Mark a `dict[int, float]` field (mapping dimension → weight) as sparse with `ZvecVectorDef(sparse=True)`: +Use the shared `SparseVector` type. Sparse values use float32. `SparseVectorSchema` is optional; zvec does not require a dimension, so `size` is accepted but ignored. ```python +from dataclasses import dataclass +from typing import Annotated + +from cocoindex.resources.schema import SparseVector, SparseVectorSchema + @dataclass -class Doc: +class SparseDoc: id: str - sparse: Annotated[dict[int, float], ZvecVectorDef(sparse=True)] + sparse: Annotated[ + SparseVector, + SparseVectorSchema(size=50_000), + ] + +schema = await zvec.CollectionSchema.from_class(SparseDoc, primary_key=["id"]) ``` +At declaration time, the value may also be an integer-to-float mapping; CocoIndex validates and sorts it before writing. The legacy `Annotated[dict[int, float], ZvecVectorDef(sparse=True)]` spelling remains supported, but `SparseVector` is preferred for portable pipeline code. `ZvecVectorDef` can still accompany the shared type when a connector-specific metric is needed. + ## Full example ```python diff --git a/python/cocoindex/_internal/datatype.py b/python/cocoindex/_internal/datatype.py index 8187c876e..81d7bffc3 100644 --- a/python/cocoindex/_internal/datatype.py +++ b/python/cocoindex/_internal/datatype.py @@ -252,13 +252,15 @@ def analyze_type_info(t: Any, *, nullable: bool = False) -> DataTypeInfo: variant = MappingType(key_type=key_type, value_type=elem_type) elif base_type in (types.UnionType, typing.Union): non_none_types = [arg for arg in type_args if arg not in (None, types.NoneType)] - if len(non_none_types) == 0: - return analyze_type_info(None) - - if len(non_none_types) == 1: - return analyze_type_info( - non_none_types[0], - nullable=nullable or len(non_none_types) < len(type_args), + if len(non_none_types) <= 1: + nested_info = analyze_type_info( + non_none_types[0] if non_none_types else None, + nullable=nullable + or bool(non_none_types) + and len(non_none_types) < len(type_args), + ) + return nested_info._replace( + annotations=nested_info.annotations + annotations ) variant = UnionType(variant_types=non_none_types) diff --git a/python/cocoindex/connectorkits/__init__.py b/python/cocoindex/connectorkits/__init__.py index 8b99bffc6..fa3466dfa 100644 --- a/python/cocoindex/connectorkits/__init__.py +++ b/python/cocoindex/connectorkits/__init__.py @@ -4,7 +4,19 @@ from typing import Any -__all__ = ["SingleWatcherGuard", "default_subpath_name"] +from .vector_schema import ( + VectorSchemas, + reject_sparse_vectors, + resolve_vector_schemas, +) + +__all__ = [ + "SingleWatcherGuard", + "VectorSchemas", + "default_subpath_name", + "reject_sparse_vectors", + "resolve_vector_schemas", +] class SingleWatcherGuard: diff --git a/python/cocoindex/connectorkits/vector_schema.py b/python/cocoindex/connectorkits/vector_schema.py new file mode 100644 index 000000000..24737ecb0 --- /dev/null +++ b/python/cocoindex/connectorkits/vector_schema.py @@ -0,0 +1,123 @@ +"""Shared vector-schema resolution for target connector implementations.""" + +from __future__ import annotations + +import collections.abc as _collections_abc +import typing as _typing + +import cocoindex as _coco +from cocoindex.resources import schema as _schema + +__all__ = ["VectorSchemas", "reject_sparse_vectors", "resolve_vector_schemas"] + + +class VectorSchemas(_typing.NamedTuple): + """Dense and sparse vector schemas resolved for one column.""" + + vector: _schema.VectorSchema | None + sparse: _schema.SparseVectorSchema | None + + +async def resolve_vector_schemas( + base_type: object, + annotations: _typing.Iterable[object], + *, + reject_sparse_vectors_for: str | None = None, +) -> VectorSchemas: + """Resolve dense and sparse metadata for a vector-aware connector. + + All direct annotations are inspected first. The first provider of each kind + is retained, but providing both dense and sparse schemas is an error. Context + keys are consulted in annotation order only when no direct schema exists, + and resolution stops at the first key that yields either schema. + + Set ``reject_sparse_vectors_for`` to the connector name when dense vectors + are supported but sparse vectors are not. Connectors that do not consume + vector metadata should call :func:`reject_sparse_vectors` instead so they do + not resolve unrelated context keys. + """ + if reject_sparse_vectors_for is not None and base_type is _schema.SparseVector: + _raise_sparse_vectors_unsupported(reject_sparse_vectors_for) + + direct_annotations: list[object] = [] + context_keys: list[_coco.ContextKey[object]] = [] + for annotation in annotations: + if isinstance(annotation, _coco.ContextKey): + context_keys.append(annotation) + else: + direct_annotations.append(annotation) + + schemas = await _resolve_providers(direct_annotations) + if schemas.vector is None and schemas.sparse is None: + for context_key in context_keys: + schemas = await _resolve_providers([_coco.use_context(context_key)]) + if schemas.vector is not None or schemas.sparse is not None: + break + + _validate_schema_base_types(base_type, schemas) + return schemas + + +def reject_sparse_vectors( + base_type: object, + annotations: _typing.Iterable[object], + *, + connector_name: str, +) -> None: + """Reject direct sparse metadata without resolving ``ContextKey`` values.""" + has_sparse_schema = any( + isinstance(annotation, _schema.SparseVectorSchemaProvider) + for annotation in annotations + ) + _validate_sparse_schema_base_type(base_type, has_sparse_schema) + if base_type is _schema.SparseVector: + _raise_sparse_vectors_unsupported(connector_name) + + +def _validate_schema_base_types(base_type: object, schemas: VectorSchemas) -> None: + _validate_sparse_schema_base_type(base_type, schemas.sparse is not None) + if schemas.vector is not None and ( + base_type is _schema.SparseVector + or ( + isinstance(base_type, type) + and issubclass(base_type, _collections_abc.Mapping) + ) + ): + raise TypeError( + f"VectorSchema requires a dense vector field, got {base_type!r}." + ) + + +def _validate_sparse_schema_base_type( + base_type: object, has_sparse_schema: bool +) -> None: + if has_sparse_schema and base_type is not _schema.SparseVector: + raise TypeError( + f"SparseVectorSchema requires a SparseVector field, got {base_type!r}." + ) + + +def _raise_sparse_vectors_unsupported(connector_name: str) -> _typing.NoReturn: + raise ValueError(f"{connector_name} does not support sparse vector columns.") + + +async def _resolve_providers(annotations: _typing.Iterable[object]) -> VectorSchemas: + vector_schema: _schema.VectorSchema | None = None + sparse_vector_schema: _schema.SparseVectorSchema | None = None + + for annotation in annotations: + if vector_schema is None and isinstance( + annotation, _schema.VectorSchemaProvider + ): + vector_schema = await annotation.__coco_vector_schema__() + if sparse_vector_schema is None and isinstance( + annotation, _schema.SparseVectorSchemaProvider + ): + sparse_vector_schema = await annotation.__coco_sparse_vector_schema__() + + if vector_schema is not None and sparse_vector_schema is not None: + raise ValueError( + "A field cannot provide both VectorSchema and SparseVectorSchema" + ) + + return VectorSchemas(vector=vector_schema, sparse=sparse_vector_schema) diff --git a/python/cocoindex/connectors/bigquery/_target.py b/python/cocoindex/connectors/bigquery/_target.py index 31c0d8109..b9522c05a 100644 --- a/python/cocoindex/connectors/bigquery/_target.py +++ b/python/cocoindex/connectors/bigquery/_target.py @@ -41,7 +41,7 @@ analyze_type_info, is_record_type, ) -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import reject_sparse_vectors, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object _RowKey = tuple[Any, ...] @@ -170,6 +170,12 @@ async def _columns_from_record_type( all_annotations.append(override) all_annotations.extend(type_info.annotations) + reject_sparse_vectors( + type_info.base_type, + all_annotations, + connector_name="BigQuery", + ) + bigquery_type_annotation = next( (t for t in all_annotations if isinstance(t, BigQueryType)), None ) diff --git a/python/cocoindex/connectors/doris/_target.py b/python/cocoindex/connectors/doris/_target.py index a62729d02..912001cf3 100644 --- a/python/cocoindex/connectors/doris/_target.py +++ b/python/cocoindex/connectors/doris/_target.py @@ -46,7 +46,7 @@ from typing_extensions import TypeVar import cocoindex as coco -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex._internal.datatype import ( AnyType, @@ -394,14 +394,12 @@ async def _columns_from_record_type( doris_type_annotation = next( (t for t in all_annotations if isinstance(t, DorisType)), None ) - vector_schema = await anext( - ( - s - for annot in all_annotations - if (s := await res_schema.get_vector_schema(annot)) is not None - ), - None, + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, + reject_sparse_vectors_for="Doris", ) + vector_schema = vector_schemas.vector if doris_type_annotation is not None: mapping = _TypeMapping( diff --git a/python/cocoindex/connectors/falkordb/_target.py b/python/cocoindex/connectors/falkordb/_target.py index 813d6377b..62db06dbc 100644 --- a/python/cocoindex/connectors/falkordb/_target.py +++ b/python/cocoindex/connectors/falkordb/_target.py @@ -17,7 +17,6 @@ import datetime import decimal import logging -import re import uuid as uuid_mod from dataclasses import dataclass from typing import ( @@ -34,7 +33,6 @@ from typing_extensions import TypeVar try: - import falkordb as _falkordb # type: ignore[import-untyped] import falkordb.asyncio as _falkordb_asyncio # type: ignore[import-untyped] except ImportError as e: raise ImportError( @@ -53,7 +51,7 @@ import msgspec import cocoindex as coco -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex._internal.datatype import ( AnyType, @@ -393,12 +391,12 @@ async def _columns_from_record_type( falkor_type_annotation = next( (t for t in all_annotations if isinstance(t, FalkorType)), None ) - vector_schema = None - for annot in all_annotations: - vs = await res_schema.get_vector_schema(annot) - if vs is not None: - vector_schema = vs - break + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, + reject_sparse_vectors_for="FalkorDB", + ) + vector_schema = vector_schemas.vector if falkor_type_annotation is not None: type_mapping = _TypeMapping( diff --git a/python/cocoindex/connectors/lancedb/_target.py b/python/cocoindex/connectors/lancedb/_target.py index 4a6000f86..87f517d56 100644 --- a/python/cocoindex/connectors/lancedb/_target.py +++ b/python/cocoindex/connectors/lancedb/_target.py @@ -38,7 +38,7 @@ import numpy as np import cocoindex as coco -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex._internal.datatype import ( AnyType, @@ -148,7 +148,9 @@ class _TypeMapping(NamedTuple): async def _get_type_mapping( - python_type: Any, *, vector_schema: res_schema.VectorSchema | None = None + python_type: Any, + *, + vector_schema: res_schema.VectorSchema | None = None, ) -> _TypeMapping: """ Get the PyArrow type mapping for a Python type. @@ -157,14 +159,13 @@ async def _get_type_mapping( Use `LanceType` annotation with `typing.Annotated` to override the default. """ type_info = analyze_type_info(python_type) + base_type = type_info.base_type # Check for LanceType annotation override for annotation in type_info.annotations: if isinstance(annotation, LanceType): return _TypeMapping(annotation.pa_type, annotation.encoder) - base_type = type_info.base_type - # Check direct leaf type mappings if base_type in _LEAF_TYPE_MAPPINGS: return _LEAF_TYPE_MAPPINGS[base_type] @@ -258,7 +259,10 @@ async def from_class( record_type: type[RowT], primary_key: list[str], *, - column_specs: dict[str, LanceType | res_schema.VectorSchemaProvider] + column_specs: dict[ + str, + LanceType | res_schema.VectorSchemaProvider, + ] | None = None, ) -> "TableSchema[RowT]": """ @@ -283,7 +287,11 @@ async def from_class( @staticmethod async def _columns_from_record_type( record_type: type, - column_specs: dict[str, LanceType | res_schema.VectorSchemaProvider] | None, + column_specs: dict[ + str, + LanceType | res_schema.VectorSchemaProvider, + ] + | None, ) -> dict[str, ColumnDef]: """Convert a record type to a dict of column name -> ColumnDef.""" record_info = RecordType(record_type) @@ -302,14 +310,12 @@ async def _columns_from_record_type( lance_type_annotation = next( (t for t in all_annotations if isinstance(t, LanceType)), None ) - vector_schema = await anext( - ( - s - for annot in all_annotations - if (s := await res_schema.get_vector_schema(annot)) is not None - ), - None, + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, + reject_sparse_vectors_for="LanceDB", ) + vector_schema = vector_schemas.vector # Determine type mapping if lance_type_annotation is not None: @@ -318,7 +324,8 @@ async def _columns_from_record_type( ) else: type_mapping = await _get_type_mapping( - field.type_hint, vector_schema=vector_schema + field.type_hint, + vector_schema=vector_schema, ) columns[field.name] = ColumnDef( diff --git a/python/cocoindex/connectors/neo4j/_target.py b/python/cocoindex/connectors/neo4j/_target.py index 862416f00..f158dbdc2 100644 --- a/python/cocoindex/connectors/neo4j/_target.py +++ b/python/cocoindex/connectors/neo4j/_target.py @@ -63,7 +63,7 @@ analyze_type_info, is_record_type, ) -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex.resources import schema as res_schema @@ -431,12 +431,12 @@ async def _columns_from_record_type( neo4j_type_annotation = next( (t for t in all_annotations if isinstance(t, Neo4jType)), None ) - vector_schema = None - for annot in all_annotations: - vs = await res_schema.get_vector_schema(annot) - if vs is not None: - vector_schema = vs - break + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, + reject_sparse_vectors_for="Neo4j", + ) + vector_schema = vector_schemas.vector if neo4j_type_annotation is not None: type_mapping = _TypeMapping( diff --git a/python/cocoindex/connectors/postgres/_target.py b/python/cocoindex/connectors/postgres/_target.py index 62590610a..dbc1862aa 100644 --- a/python/cocoindex/connectors/postgres/_target.py +++ b/python/cocoindex/connectors/postgres/_target.py @@ -41,7 +41,7 @@ import numpy as np import cocoindex as coco -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex._internal.datatype import ( AnyType, @@ -173,7 +173,27 @@ def _vector_encoder(value: Any) -> str: return "[" + ",".join(str(float(x)) for x in value) + "]" -_PGVECTOR_TYPE_BASES: frozenset[str] = frozenset({"vector", "halfvec"}) +def _make_sparsevec_encoder(dim: int) -> ValueEncoder: + """Build an encoder for pgvector's 1-based sparsevec text format.""" + + def encode(value: Any) -> str: + sparse_vector = res_schema.as_sparse_vector(value) + if sparse_vector.indices and sparse_vector.indices[-1] >= dim: + raise ValueError( + f"sparse vector index {sparse_vector.indices[-1]} out of range " + f"for sparsevec({dim})" + ) + elems = ",".join( + f"{index + 1}:{float(vector_value)}" + for index, vector_value in zip(sparse_vector.indices, sparse_vector.values) + if vector_value != 0.0 + ) + return "{" + elems + "}/" + str(dim) + + return encode + + +_PGVECTOR_TYPE_BASES: frozenset[str] = frozenset({"vector", "halfvec", "sparsevec"}) def _pgvector_type_base(pg_type: str) -> str | None: @@ -254,10 +274,16 @@ class _TypeMapping(NamedTuple): # Default mapping for complex types that need JSON encoding _JSONB_MAPPING = _TypeMapping("jsonb", _json_encoder) +_SPARSEVEC_DIMENSION_ERROR = ( + "pgvector sparsevec requires a dimension; provide SparseVectorSchema(size=...)" +) async def _get_type_mapping( - python_type: Any, *, vector_schema: res_schema.VectorSchema | None = None + python_type: Any, + *, + vector_schema: res_schema.VectorSchema | None = None, + sparse_vector_schema: res_schema.SparseVectorSchema | None = None, ) -> _TypeMapping: """ Get the PostgreSQL type mapping for a Python type. @@ -266,16 +292,20 @@ async def _get_type_mapping( https://magicstack.github.io/asyncpg/current/usage.html#type-conversion For types that map to multiple PostgreSQL types, uses the broader one. - Use `PgType` annotation with `typing.Annotated` to override the default. """ type_info = analyze_type_info(python_type) + base_type = type_info.base_type - # Check for PgType annotation override - for annotation in type_info.annotations: - if isinstance(annotation, PgType): - return _TypeMapping(annotation.pg_type, annotation.encoder) + if sparse_vector_schema is not None: + if sparse_vector_schema.size is None: + raise ValueError(_SPARSEVEC_DIMENSION_ERROR) + return _TypeMapping( + f"sparsevec({sparse_vector_schema.size})", + _make_sparsevec_encoder(sparse_vector_schema.size), + ) - base_type = type_info.base_type + if base_type is res_schema.SparseVector: + raise ValueError(_SPARSEVEC_DIMENSION_ERROR) # Check direct leaf type mappings if base_type in _LEAF_TYPE_MAPPINGS: @@ -366,7 +396,12 @@ async def from_class( record_type: type[RowT], primary_key: list[str], *, - column_overrides: dict[str, PgType | res_schema.VectorSchemaProvider] + column_overrides: dict[ + str, + PgType + | res_schema.VectorSchemaProvider + | res_schema.SparseVectorSchemaProvider, + ] | None = None, ) -> "TableSchema[RowT]": """ @@ -379,7 +414,7 @@ async def from_class( record_type: A record type (dataclass, NamedTuple, or Pydantic model). primary_key: List of column names that form the primary key. column_overrides: Optional dict mapping column names to PgType or - VectorSchemaProvider to override the default type mapping. + vector schema provider to override the default type mapping. """ if not is_record_type(record_type): raise TypeError( @@ -392,7 +427,13 @@ async def from_class( @staticmethod async def _columns_from_record_type( record_type: type, - column_overrides: dict[str, PgType | res_schema.VectorSchemaProvider] | None, + column_overrides: dict[ + str, + PgType + | res_schema.VectorSchemaProvider + | res_schema.SparseVectorSchemaProvider, + ] + | None, ) -> dict[str, ColumnDef]: """Convert a record type to a dict of column name -> ColumnDef.""" record_info = RecordType(record_type) @@ -412,23 +453,31 @@ async def _columns_from_record_type( pg_type_annotation = next( (t for t in all_annotations if isinstance(t, PgType)), None ) - vector_schema = await anext( - ( - s - for annot in all_annotations - if (s := await res_schema.get_vector_schema(annot)) is not None - ), - None, + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, ) + vector_schema = vector_schemas.vector + sparse_vector_schema = vector_schemas.sparse # Determine type mapping if pg_type_annotation is not None: + if ( + sparse_vector_schema is not None + or type_info.base_type is res_schema.SparseVector + ): + raise ValueError( + f"Column {field.name!r} cannot combine PgType with " + "sparse vector metadata." + ) type_mapping = _TypeMapping( pg_type_annotation.pg_type, pg_type_annotation.encoder ) else: type_mapping = await _get_type_mapping( - field.type_hint, vector_schema=vector_schema + field.type_hint, + vector_schema=vector_schema, + sparse_vector_schema=sparse_vector_schema, ) columns[field.name] = ColumnDef( @@ -460,6 +509,11 @@ class _RowAction(NamedTuple): "l2": "halfvec_l2_ops", "ip": "halfvec_ip_ops", }, + "sparsevec": { + "cosine": "sparsevec_cosine_ops", + "l2": "sparsevec_l2_ops", + "ip": "sparsevec_ip_ops", + }, } @@ -1332,6 +1386,9 @@ def declare_vector_index( lists: Number of lists (ivfflat only). m: Maximum number of connections per layer (hnsw only). ef_construction: Size of the dynamic candidate list (hnsw only). + + Sparse-vector columns only support HNSW indexes. pgvector limits + HNSW-indexed sparse vectors to 1,000 non-zero elements. """ if name is None: name = column @@ -1340,6 +1397,10 @@ def declare_vector_index( raise ValueError( f"Column '{column}' not found in table schema: {list(self._table_schema.columns.keys())}" ) + if _pgvector_type_base(col_def.type) == "sparsevec" and method != "hnsw": + raise ValueError( + 'pgvector only supports HNSW indexes on sparsevec columns; pass method="hnsw"' + ) spec = _VectorIndexSpec( column=column, metric=metric, diff --git a/python/cocoindex/connectors/qdrant/_target.py b/python/cocoindex/connectors/qdrant/_target.py index 4ada1d3e9..a3c3e9636 100644 --- a/python/cocoindex/connectors/qdrant/_target.py +++ b/python/cocoindex/connectors/qdrant/_target.py @@ -18,6 +18,7 @@ Collection, Generic, Literal, + Mapping, NamedTuple, Sequence, cast, @@ -42,6 +43,17 @@ # Public alias for Qdrant point model PointStruct = qdrant_models.PointStruct + +def sparse_vector( + value: res_schema.SparseVector | Mapping[int, float], +) -> qdrant_models.SparseVector: + """Convert a canonical sparse vector value to Qdrant's model.""" + canonical = res_schema.as_sparse_vector(value) + return qdrant_models.SparseVector( + indices=list(canonical.indices), values=list(canonical.values) + ) + + # Type aliases _PointId = str | int _PointFingerprint = bytes @@ -775,4 +787,5 @@ def _validate_point_id(raw: object) -> _PointId: "create_client", "declare_collection_target", "mount_collection_target", + "sparse_vector", ] diff --git a/python/cocoindex/connectors/snowflake/_target.py b/python/cocoindex/connectors/snowflake/_target.py index b136a64fa..322b9cad6 100644 --- a/python/cocoindex/connectors/snowflake/_target.py +++ b/python/cocoindex/connectors/snowflake/_target.py @@ -41,7 +41,7 @@ analyze_type_info, is_record_type, ) -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import reject_sparse_vectors, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object _RowKey = tuple[Any, ...] @@ -163,6 +163,12 @@ async def _columns_from_record_type( all_annotations.append(override) all_annotations.extend(type_info.annotations) + reject_sparse_vectors( + type_info.base_type, + all_annotations, + connector_name="Snowflake", + ) + snowflake_type_annotation = next( (t for t in all_annotations if isinstance(t, SnowflakeType)), None ) diff --git a/python/cocoindex/connectors/sqlite/_target.py b/python/cocoindex/connectors/sqlite/_target.py index 33d017950..038e958ae 100644 --- a/python/cocoindex/connectors/sqlite/_target.py +++ b/python/cocoindex/connectors/sqlite/_target.py @@ -35,7 +35,7 @@ import numpy as np import cocoindex as coco -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex._internal.rwlock import RWLock import msgspec @@ -226,7 +226,9 @@ class _TypeMapping(NamedTuple): async def _get_type_mapping( - python_type: Any, *, vector_schema: res_schema.VectorSchema | None = None + python_type: Any, + *, + vector_schema: res_schema.VectorSchema | None = None, ) -> _TypeMapping: """ Get the SQLite type mapping for a Python type. @@ -240,14 +242,13 @@ async def _get_type_mapping( Use `SqliteType` annotation with `typing.Annotated` to override the default. """ type_info = analyze_type_info(python_type) + base_type = type_info.base_type # Check for SqliteType annotation override for annotation in type_info.annotations: if isinstance(annotation, SqliteType): return _TypeMapping(annotation.sqlite_type, annotation.encoder) - base_type = type_info.base_type - # Check direct leaf type mappings if base_type in _LEAF_TYPE_MAPPINGS: return _LEAF_TYPE_MAPPINGS[base_type] @@ -340,7 +341,10 @@ async def from_class( record_type: type[RowT], primary_key: list[str], *, - column_overrides: dict[str, SqliteType | res_schema.VectorSchemaProvider] + column_overrides: dict[ + str, + SqliteType | res_schema.VectorSchemaProvider, + ] | None = None, ) -> "TableSchema[RowT]": """ @@ -365,16 +369,21 @@ async def from_class( @staticmethod async def _columns_from_record_type( record_type: type, - column_overrides: dict[str, SqliteType | res_schema.VectorSchemaProvider] + column_overrides: dict[ + str, + SqliteType | res_schema.VectorSchemaProvider, + ] | None, ) -> dict[str, ColumnDef]: """Convert a record type to a dict of column name -> ColumnDef.""" record_info = RecordType(record_type) columns: dict[str, ColumnDef] = {} - for field in record_info.fields: - override = column_overrides.get(field.name) if column_overrides else None - type_info = analyze_type_info(field.type_hint) + for record_field in record_info.fields: + override = ( + column_overrides.get(record_field.name) if column_overrides else None + ) + type_info = analyze_type_info(record_field.type_hint) all_annotations = [] if override is not None: @@ -385,14 +394,12 @@ async def _columns_from_record_type( sqlite_type_annotation = next( (t for t in all_annotations if isinstance(t, SqliteType)), None ) - vector_schema = await anext( - ( - s - for annot in all_annotations - if (s := await res_schema.get_vector_schema(annot)) is not None - ), - None, + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, + reject_sparse_vectors_for="SQLite", ) + vector_schema = vector_schemas.vector # Determine type mapping if sqlite_type_annotation is not None: @@ -401,10 +408,11 @@ async def _columns_from_record_type( ) else: type_mapping = await _get_type_mapping( - field.type_hint, vector_schema=vector_schema + record_field.type_hint, + vector_schema=vector_schema, ) - columns[field.name] = ColumnDef( + columns[record_field.name] = ColumnDef( type=type_mapping.sqlite_type.strip(), nullable=type_info.nullable, encoder=type_mapping.encoder, diff --git a/python/cocoindex/connectors/surrealdb/_target.py b/python/cocoindex/connectors/surrealdb/_target.py index 2a13094b0..9114d7b63 100644 --- a/python/cocoindex/connectors/surrealdb/_target.py +++ b/python/cocoindex/connectors/surrealdb/_target.py @@ -47,7 +47,7 @@ import numpy as np import cocoindex as coco -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex._internal.datatype import ( AnyType, @@ -407,14 +407,12 @@ async def _columns_from_record_type( surreal_type_annotation = next( (t for t in all_annotations if isinstance(t, SurrealType)), None ) - vector_schema = await anext( - ( - s - for annot in all_annotations - if (s := await res_schema.get_vector_schema(annot)) is not None - ), - None, + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + all_annotations, + reject_sparse_vectors_for="SurrealDB", ) + vector_schema = vector_schemas.vector # Determine type mapping if surreal_type_annotation is not None: diff --git a/python/cocoindex/connectors/turbopuffer/_target.py b/python/cocoindex/connectors/turbopuffer/_target.py index 638bfc626..a0f26beaf 100644 --- a/python/cocoindex/connectors/turbopuffer/_target.py +++ b/python/cocoindex/connectors/turbopuffer/_target.py @@ -12,6 +12,7 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from typing import ( Any, @@ -80,6 +81,8 @@ class _ResolvedNamedVectorsDef(msgspec.Struct, frozen=True, tag=True): async def _resolve_vector_def(vector_def: VectorDef) -> _ResolvedVectorDef: vs = await res_schema.get_vector_schema(vector_def.schema) if vs is None: + if await res_schema.get_sparse_vector_schema(vector_def.schema) is not None: + raise ValueError("Turbopuffer does not support sparse vector schemas.") raise ValueError(f"Invalid vector definition: {vector_def}") # Validate dtype upfront so bad schemas fail at construction time, not on # the first write. Discards the return — used for its raise side effect. @@ -189,6 +192,8 @@ class Row: def _vector_to_list(v: Sequence[float] | np.ndarray) -> list[float]: + if isinstance(v, (res_schema.SparseVector, Mapping)): + raise ValueError("Turbopuffer does not support sparse vector values.") if isinstance(v, np.ndarray): return v.tolist() # type: ignore[no-any-return] return list(v) diff --git a/python/cocoindex/connectors/valkey/_target.py b/python/cocoindex/connectors/valkey/_target.py index ecf392238..3ed7fac39 100644 --- a/python/cocoindex/connectors/valkey/_target.py +++ b/python/cocoindex/connectors/valkey/_target.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import collections.abc import logging import re import struct @@ -19,6 +20,13 @@ import msgspec import numpy as np +import cocoindex as coco +from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits.fingerprint import fingerprint_object +from cocoindex.resources import schema as res_schema +from cocoindex._internal.context_keys import ContextKey, ContextProvider +from cocoindex._internal.datatype import TypeChecker + logger = logging.getLogger(__name__) try: @@ -48,13 +56,6 @@ "Please install cocoindex[valkey]." ) from e -import cocoindex as coco -from cocoindex.connectorkits import statediff, target -from cocoindex.connectorkits.fingerprint import fingerprint_object -from cocoindex.resources import schema as res_schema -from cocoindex._internal.context_keys import ContextKey, ContextProvider -from cocoindex._internal.datatype import TypeChecker - # --------------------------------------------------------------------------- # Public types @@ -257,6 +258,8 @@ def _validate_name(value: str, label: str) -> str: def _vector_to_bytes(vector: list[float] | np.ndarray) -> bytes: # type: ignore[type-arg] """Pack a vector into little-endian float32 bytes for Valkey HASH storage.""" + if isinstance(vector, (collections.abc.Mapping, res_schema.SparseVector)): + raise ValueError("Valkey does not support sparse vector values.") if isinstance(vector, np.ndarray): return vector.astype(np.float32).tobytes() return struct.pack(f"<{len(vector)}f", *vector) diff --git a/python/cocoindex/connectors/zvec/_target.py b/python/cocoindex/connectors/zvec/_target.py index 2e1012cce..0c5f668ae 100644 --- a/python/cocoindex/connectors/zvec/_target.py +++ b/python/cocoindex/connectors/zvec/_target.py @@ -22,6 +22,7 @@ import re import threading import uuid +from collections.abc import Mapping from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path @@ -50,10 +51,11 @@ import msgspec import cocoindex as coco -from cocoindex.connectorkits import statediff, target +from cocoindex.connectorkits import resolve_vector_schemas, statediff, target from cocoindex.connectorkits.fingerprint import fingerprint_object from cocoindex._internal.context_keys import ContextKey, ContextProvider from cocoindex._internal.datatype import ( + MappingType, RecordType, SequenceType, TypeChecker, @@ -233,7 +235,9 @@ class ZvecVectorDef(NamedTuple): ``VectorSchema`` (via ``Annotated`` or ``column_overrides``). This annotation tunes the index and marks sparse fields. - For sparse vectors, set ``sparse=True`` on a ``dict[int, float]`` field. + For sparse vectors, prefer a ``SparseVector`` field, optionally annotated + with ``SparseVectorSchema``. ``sparse=True`` on a ``dict[int, float]`` field + remains supported for compatibility. """ metric: Literal["cosine", "ip", "l2"] = "cosine" @@ -342,6 +346,7 @@ async def _resolve_column( | ZvecVectorDef | ZvecFtsType | res_schema.VectorSchemaProvider + | res_schema.SparseVectorSchemaProvider | None, ) -> _Column: type_info = analyze_type_info(type_hint) @@ -351,18 +356,18 @@ async def _resolve_column( annotations.append(override) annotations.extend(type_info.annotations) - vector_schema: res_schema.VectorSchema | None = None - for annot in annotations: - vs = await res_schema.get_vector_schema(annot) - if vs is not None: - vector_schema = vs - break + vector_schemas = await resolve_vector_schemas( + type_info.base_type, + annotations, + ) + vector_schema = vector_schemas.vector + sparse_vector_schema = vector_schemas.sparse vector_def = next((a for a in annotations if isinstance(a, ZvecVectorDef)), None) zvec_type = next((a for a in annotations if isinstance(a, ZvecType)), None) fts_type = next((a for a in annotations if isinstance(a, ZvecFtsType)), None) - # Dense vector: NumPy ndarray with a VectorSchema. + # Dense vector: a VectorSchema marks any sequence accepted by zvec. if vector_schema is not None: if vector_schema.size <= 0: raise ValueError( @@ -379,14 +384,33 @@ async def _resolve_column( quantize=vd.quantize, ) - # Sparse vector: explicitly marked via ZvecVectorDef(sparse=True). - if vector_def is not None and vector_def.sparse: + # Sparse vector: shared canonical type/schema or legacy sparse vector def. + if ( + sparse_vector_schema is not None + or (vector_def is not None and vector_def.sparse) + or type_info.base_type is res_schema.SparseVector + ): + if ( + vector_def is not None + and vector_def.sparse + and type_info.base_type is not res_schema.SparseVector + and not isinstance(type_info.variant, MappingType) + ): + raise ValueError( + f"ZvecVectorDef(sparse=True) on column {name!r} requires a " + f"SparseVector or mapping field, got {type_info.base_type!r}." + ) + if zvec_type is not None: + raise ValueError( + f"Column {name!r} cannot combine ZvecType with sparse vector metadata." + ) + vd = vector_def or ZvecVectorDef() return _Column( name=name, kind="sparse", data_type=_zvec.DataType.SPARSE_VECTOR_FP32, nullable=type_info.nullable, - metric=vector_def.metric, + metric=vd.metric, ) if type_info.base_type is np.ndarray: @@ -480,7 +504,11 @@ async def from_class( *, column_overrides: dict[ str, - ZvecType | ZvecVectorDef | ZvecFtsType | res_schema.VectorSchemaProvider, + ZvecType + | ZvecVectorDef + | ZvecFtsType + | res_schema.VectorSchemaProvider + | res_schema.SparseVectorSchemaProvider, ] | None = None, ) -> "CollectionSchema[RowT]": @@ -870,11 +898,18 @@ def _row_get(row: Any, name: str) -> Any: def _to_float_list(value: Any) -> list[float]: + if isinstance(value, (Mapping, res_schema.SparseVector)): + raise ValueError("zvec does not support sparse values in dense vector columns.") if isinstance(value, np.ndarray): return cast(list[float], value.astype(float).tolist()) return [float(x) for x in value] +def _to_sparse_dict(value: Any) -> dict[int, float]: + sparse_vector = res_schema.as_sparse_vector(value) + return dict(zip(sparse_vector.indices, sparse_vector.values)) + + class CollectionTarget( Generic[RowT, coco.MaybePendingS], coco.ResolvesTo["CollectionTarget[RowT]"] ): @@ -913,11 +948,10 @@ def declare_row(self: "CollectionTarget[RowT]", *, row: RowT) -> None: if col.kind == "dense": vectors[name] = None if value is None else _to_float_list(value) elif col.kind == "sparse": - vectors[name] = ( - None - if value is None - else {int(k): float(v) for k, v in dict(value).items()} - ) + if value is None: + vectors[name] = None + else: + vectors[name] = _to_sparse_dict(value) else: if value is not None and col.encoder is not None: value = col.encoder(value) diff --git a/python/cocoindex/resources/schema.py b/python/cocoindex/resources/schema.py index 975814484..e77d82e30 100644 --- a/python/cocoindex/resources/schema.py +++ b/python/cocoindex/resources/schema.py @@ -7,12 +7,164 @@ from __future__ import annotations +import collections.abc as _collections_abc +import math as _math +import numbers as _numbers +import operator as _operator import typing as _typing -import cocoindex as coco +import cocoindex as _coco import msgspec as _msgspec import numpy as _np +class SparseVector(_msgspec.Struct, frozen=True): + """Canonical sparse vector value. + + Indices are 0-based, sorted ascending, and unique. Values are stored in a + parallel tuple so the representation is deterministic and serializable. + """ + + indices: tuple[int, ...] + values: tuple[float, ...] + + def __post_init__(self) -> None: + if not isinstance(self.indices, tuple): + raise TypeError("indices must be a tuple of integers") + if not isinstance(self.values, tuple): + raise TypeError("values must be a tuple of floats") + if len(self.indices) != len(self.values): + raise ValueError("indices and values must have the same length") + previous_index = -1 + for index, value in zip(self.indices, self.values, strict=True): + if type(index) is not int: + raise TypeError("indices must contain only integers") + if type(value) is not float: + raise TypeError("values must contain only floats") + if index < 0: + raise ValueError("indices must be non-negative") + if index <= previous_index: + raise ValueError("indices must be sorted ascending and unique") + if not _math.isfinite(value): + raise ValueError("values must be finite") + previous_index = index + + @classmethod + def from_arrays( + cls, + indices: _typing.Iterable[_typing.SupportsIndex], + values: _typing.Iterable[_typing.SupportsFloat], + ) -> SparseVector: + """Coerce parallel numeric arrays and sort them by ascending index.""" + normalized_indices = _normalize_sparse_indices(indices) + normalized_values = _normalize_sparse_values(values) + if len(normalized_indices) != len(normalized_values): + raise ValueError("indices and values must have the same length") + items = sorted(zip(normalized_indices, normalized_values, strict=True)) + return cls( + indices=tuple(index for index, _ in items), + values=tuple(value for _, value in items), + ) + + @classmethod + def from_mapping(cls, m: _typing.Mapping[int, float]) -> SparseVector: + """Coerce and sort an index-to-value mapping.""" + if not isinstance(m, _collections_abc.Mapping): + raise TypeError("expected a Mapping[int, float]") + return cls.from_arrays(m.keys(), m.values()) + + +def _normalize_sparse_indices(indices: object) -> tuple[int, ...]: + try: + iterator = iter(indices) # type: ignore[call-overload] + except TypeError as e: + raise TypeError("indices must be an iterable of integers") from e + + normalized: list[int] = [] + for index in iterator: + if isinstance(index, (bool, _np.bool_)): + raise TypeError("sparse vector indices must be integers, not bool") + try: + normalized.append(_operator.index(index)) + except TypeError as e: + raise TypeError( + f"sparse vector indices must be integers, got {index!r}" + ) from e + return tuple(normalized) + + +def _normalize_sparse_values(values: object) -> tuple[float, ...]: + try: + iterator = iter(values) # type: ignore[call-overload] + except TypeError as e: + raise TypeError("values must be an iterable of real numbers") from e + + normalized: list[float] = [] + for value in iterator: + value_type = type(value) + if value_type is float: + normalized.append(value) + continue + if value_type is int: + normalized.append(float(value)) + continue + if isinstance(value, (bool, _np.bool_)) or not isinstance(value, _numbers.Real): + raise TypeError(f"sparse vector values must be real numbers, got {value!r}") + normalized.append(float(value)) + return tuple(normalized) + + +def as_sparse_vector( + v: SparseVector | _typing.Mapping[int, float], +) -> SparseVector: + """Normalize a sparse vector value to the canonical representation.""" + if isinstance(v, SparseVector): + return v + if not isinstance(v, _collections_abc.Mapping): + raise TypeError( + f"expected SparseVector or Mapping[int, float], got {type(v).__name__}" + ) + return SparseVector.from_mapping(v) + + +@_typing.runtime_checkable +class SparseVectorSchemaProvider(_typing.Protocol): + """Provider of additional information for a sparse vector column.""" + + def __coco_sparse_vector_schema__( + self, + ) -> _typing.Awaitable[SparseVectorSchema]: ... + + +class SparseVectorSchema(_msgspec.Struct, frozen=True, tag=True): + """Additional information for a sparse vector column.""" + + size: int | None = None + + def __post_init__(self) -> None: + if self.size is not None: + if isinstance(self.size, (bool, _np.bool_)): + raise TypeError("sparse vector size must be an integer or None") + try: + size = _operator.index(self.size) + except TypeError as e: + raise TypeError("sparse vector size must be an integer or None") from e + if size <= 0: + raise ValueError("sparse vector size must be positive") + _msgspec.structs.force_setattr(self, "size", size) + + async def __coco_sparse_vector_schema__(self) -> SparseVectorSchema: + return self + + +async def get_sparse_vector_schema(obj: object) -> SparseVectorSchema | None: + """Get the sparse vector schema from an object, if it provides one.""" + if isinstance(obj, _coco.ContextKey): + obj = _coco.use_context(obj) + if isinstance(obj, SparseVectorSchemaProvider): + return await obj.__coco_sparse_vector_schema__() + return None + + @_typing.runtime_checkable class VectorSchemaProvider(_typing.Protocol): """Additional information for a vector column.""" @@ -32,8 +184,8 @@ async def __coco_vector_schema__(self) -> VectorSchema: async def get_vector_schema(obj: object) -> VectorSchema | None: """Helper function to get the vector schema from an object, if it provides one.""" - if isinstance(obj, coco.ContextKey): - obj = coco.use_context(obj) + if isinstance(obj, _coco.ContextKey): + obj = _coco.use_context(obj) if isinstance(obj, VectorSchemaProvider): return await obj.__coco_vector_schema__() return None @@ -57,8 +209,8 @@ async def __coco_multi_vector_schema__(self) -> MultiVectorSchema: async def get_multi_vector_schema(obj: object) -> MultiVectorSchema | None: """Helper function to get the multi-vector schema from an object, if it provides one.""" - if isinstance(obj, coco.ContextKey): - obj = coco.use_context(obj) + if isinstance(obj, _coco.ContextKey): + obj = _coco.use_context(obj) if isinstance(obj, MultiVectorSchemaProvider): return await obj.__coco_multi_vector_schema__() return None @@ -67,8 +219,13 @@ async def get_multi_vector_schema(obj: object) -> MultiVectorSchema | None: __all__ = [ "MultiVectorSchema", "MultiVectorSchemaProvider", + "SparseVector", + "SparseVectorSchema", + "SparseVectorSchemaProvider", "VectorSchema", "VectorSchemaProvider", + "as_sparse_vector", "get_multi_vector_schema", + "get_sparse_vector_schema", "get_vector_schema", ] diff --git a/python/tests/connectorkits/__init__.py b/python/tests/connectorkits/__init__.py new file mode 100644 index 000000000..73da7409d --- /dev/null +++ b/python/tests/connectorkits/__init__.py @@ -0,0 +1 @@ +"""Tests for shared connector implementation helpers.""" diff --git a/python/tests/connectorkits/test_vector_schema.py b/python/tests/connectorkits/test_vector_schema.py new file mode 100644 index 000000000..b18907244 --- /dev/null +++ b/python/tests/connectorkits/test_vector_schema.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +import collections.abc +from typing import Annotated + +import numpy as np +import pytest +from numpy.typing import NDArray + +import cocoindex as coco +from cocoindex._internal.datatype import analyze_type_info +from cocoindex.connectorkits import reject_sparse_vectors, resolve_vector_schemas +from cocoindex.resources.schema import ( + SparseVector, + SparseVectorSchema, + VectorSchema, +) + +from tests import common + +_DENSE_SCHEMA = VectorSchema(dtype=np.dtype(np.float32), size=4) +_DENSE_SCHEMA_KEY = coco.ContextKey[VectorSchema]( + "test_connectorkits_vector_schema/dense" +) +_UNRELATED_CONTEXT = coco.ContextKey[str]("test_connectorkits_vector_schema/unrelated") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_first", [False, True]) +async def test_direct_schema_wins_without_resolving_context( + context_first: bool, +) -> None: + annotations: tuple[object, ...] = ( + (_UNRELATED_CONTEXT, _DENSE_SCHEMA) + if context_first + else (_DENSE_SCHEMA, _UNRELATED_CONTEXT) + ) + + schemas = await resolve_vector_schemas(np.ndarray, annotations) + + assert schemas.vector is _DENSE_SCHEMA + assert schemas.sparse is None + + +@pytest.mark.asyncio +async def test_context_resolution_stops_after_first_schema() -> None: + env = common.create_test_env(__file__) + env.context_provider.provide(_DENSE_SCHEMA_KEY, _DENSE_SCHEMA) + + async def resolve() -> None: + schemas = await resolve_vector_schemas( + np.ndarray, [_DENSE_SCHEMA_KEY, _UNRELATED_CONTEXT] + ) + assert schemas.vector is _DENSE_SCHEMA + assert schemas.sparse is None + + app = coco.App( + coco.AppConfig(name="test_connectorkits_context_schema", environment=env), + resolve, + ) + await app.update() + + +@pytest.mark.asyncio +async def test_resolver_rejects_dense_sparse_conflict() -> None: + with pytest.raises(ValueError, match="both VectorSchema and SparseVectorSchema"): + await resolve_vector_schemas( + SparseVector, + [_DENSE_SCHEMA, SparseVectorSchema()], + ) + + +@pytest.mark.asyncio +async def test_resolver_validates_sparse_schema_base_type() -> None: + with pytest.raises(TypeError, match="requires a SparseVector field"): + await resolve_vector_schemas(str, [SparseVectorSchema()]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("base_type", [SparseVector, dict, collections.abc.Mapping]) +async def test_resolver_rejects_dense_schema_on_sparse_shapes( + base_type: object, +) -> None: + with pytest.raises(TypeError, match="VectorSchema requires a dense vector field"): + await resolve_vector_schemas(base_type, [_DENSE_SCHEMA]) + + +@pytest.mark.asyncio +async def test_nullable_schema_resolution_preserves_inner_first_precedence() -> None: + inner_schema = VectorSchema(dtype=np.dtype(np.float32), size=384) + outer_schema = VectorSchema(dtype=np.dtype(np.float32), size=768) + field_type = Annotated[ + Annotated[NDArray[np.float32], inner_schema] | None, + outer_schema, + ] + + type_info = analyze_type_info(field_type) + schemas = await resolve_vector_schemas(type_info.base_type, type_info.annotations) + + assert type_info.nullable + assert schemas.vector is inner_schema + + +@pytest.mark.parametrize( + ("base_type", "annotations"), + [ + (SparseVector, ()), + (SparseVector, (SparseVectorSchema(),)), + ], +) +def test_non_resolving_guard_rejects_sparse_columns( + base_type: object, annotations: tuple[object, ...] +) -> None: + with pytest.raises(ValueError, match="Example does not support sparse vector"): + reject_sparse_vectors( + base_type, + annotations, + connector_name="Example", + ) + + +def test_non_resolving_guard_ignores_context_keys() -> None: + reject_sparse_vectors( + np.ndarray, + [_UNRELATED_CONTEXT], + connector_name="Example", + ) + + +@pytest.mark.asyncio +async def test_resolver_can_reject_sparse_for_dense_only_connector() -> None: + with pytest.raises(ValueError, match="Example does not support sparse vector"): + await resolve_vector_schemas( + SparseVector, + [], + reject_sparse_vectors_for="Example", + ) diff --git a/python/tests/connectors/test_bigquery_target.py b/python/tests/connectors/test_bigquery_target.py index 08969b872..d66af008b 100644 --- a/python/tests/connectors/test_bigquery_target.py +++ b/python/tests/connectors/test_bigquery_target.py @@ -9,7 +9,9 @@ import uuid from typing import Annotated, Any, cast +import numpy as np import pytest +from numpy.typing import NDArray import cocoindex as coco from cocoindex.connectorkits import target @@ -17,6 +19,9 @@ from cocoindex.connectors.bigquery import _target BIGQUERY_DB = coco.ContextKey[bigquery.ConnectionConfig]("bigquery_test_db") +_UNPROVIDED_VECTOR_SCHEMA = coco.ContextKey[object]( + "bigquery_test_unprovided_vector_schema" +) @dataclasses.dataclass @@ -44,6 +49,12 @@ class OverrideRow: vector: Annotated[list[float], bigquery.BigQueryType("ARRAY")] +@dataclasses.dataclass +class ContextAnnotatedVectorRow: + id: int + embedding: Annotated[NDArray[np.float32], _UNPROVIDED_VECTOR_SCHEMA] + + class FakeClient: def __init__(self) -> None: self.calls: list[tuple[str, tuple[_target.QueryParam, ...]]] = [] @@ -109,6 +120,15 @@ async def test_table_target_rejects_invalid_table_identifier() -> None: ) +@pytest.mark.asyncio +async def test_table_schema_does_not_resolve_vector_context_annotations() -> None: + schema = await bigquery.TableSchema.from_class( + ContextAnnotatedVectorRow, primary_key=["id"] + ) + + assert schema.columns["embedding"].type == "JSON" + + @pytest.mark.asyncio async def test_table_target_rejects_invalid_column_identifier() -> None: schema = bigquery.TableSchema( diff --git a/python/tests/connectors/test_lancedb_target.py b/python/tests/connectors/test_lancedb_target.py index 438db2e2b..8ca1ec05c 100644 --- a/python/tests/connectors/test_lancedb_target.py +++ b/python/tests/connectors/test_lancedb_target.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import Any, Iterator, NamedTuple, cast +import numpy as np import pytest import cocoindex as coco diff --git a/python/tests/connectors/test_postgres_sparse_vector.py b/python/tests/connectors/test_postgres_sparse_vector.py new file mode 100644 index 000000000..3e8770d12 --- /dev/null +++ b/python/tests/connectors/test_postgres_sparse_vector.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Annotated, Any, cast + +import pytest + +from cocoindex.resources.schema import SparseVector, SparseVectorSchema + +try: + from cocoindex.connectors import postgres + from cocoindex.connectors.postgres._target import _make_sparsevec_encoder + + _HAS_POSTGRES = True +except ImportError: + _HAS_POSTGRES = False + +pytestmark = pytest.mark.skipif( + not _HAS_POSTGRES, reason="postgres dependencies not installed" +) + + +@dataclass +class _SparseVectorRow: + id: str + embedding: Annotated[ + SparseVector, + SparseVectorSchema(size=100), + ] + + +@pytest.mark.asyncio +async def test_postgres_sparse_vector_schema_and_encoder() -> None: + schema = await postgres.TableSchema.from_class(_SparseVectorRow, primary_key=["id"]) + sparse_column = schema.columns["embedding"] + assert sparse_column.type == "sparsevec(100)" + assert sparse_column.encoder is not None + assert sparse_column.encoder({7: 0.9, 1: 0.5}) == "{2:0.5,8:0.9}/100" + + encoder = _make_sparsevec_encoder(100) + assert encoder(SparseVector(indices=(1, 7), values=(0.0, 0.9))) == "{8:0.9}/100" + assert encoder(SparseVector(indices=(), values=())) == "{}/100" + with pytest.raises(ValueError, match="index 100 out of range"): + encoder(SparseVector(indices=(100,), values=(0.5,))) + with pytest.raises(ValueError, match="values must be finite"): + encoder({1: float("nan")}) + + @dataclass + class OverrideSparseRow: + id: str + embedding: SparseVector + + override_schema = await postgres.TableSchema.from_class( + OverrideSparseRow, + primary_key=["id"], + column_overrides={"embedding": SparseVectorSchema(size=50)}, + ) + assert override_schema.columns["embedding"].type == "sparsevec(50)" + + +@pytest.mark.asyncio +async def test_postgres_nullable_sparse_vector_preserves_schema() -> None: + @dataclass + class NullableSparseRow: + id: str + embedding: Annotated[ + SparseVector | None, + SparseVectorSchema(size=100), + ] + + schema = await postgres.TableSchema.from_class( + NullableSparseRow, primary_key=["id"] + ) + + assert schema.columns["embedding"].type == "sparsevec(100)" + assert schema.columns["embedding"].nullable + + +@pytest.mark.asyncio +async def test_postgres_rejects_native_override_with_sparse_metadata() -> None: + @dataclass + class NativeAndSparseRow: + id: str + embedding: Annotated[ + SparseVector, + SparseVectorSchema(size=100), + ] + + with pytest.raises(ValueError, match="cannot combine PgType"): + await postgres.TableSchema.from_class( + NativeAndSparseRow, + primary_key=["id"], + column_overrides={"embedding": postgres.PgType("sparsevec(100)")}, + ) + + +@pytest.mark.asyncio +async def test_postgres_sparse_vector_schema_requires_dimension() -> None: + @dataclass + class MissingDimensionRow: + id: str + embedding: Annotated[ + SparseVector, + SparseVectorSchema(), + ] + + @dataclass + class MissingAnnotationRow: + id: str + embedding: SparseVector + + for row_type in (MissingDimensionRow, MissingAnnotationRow): + with pytest.raises(ValueError, match="sparsevec requires a dimension.*size"): + await postgres.TableSchema.from_class(row_type, primary_key=["id"]) + + +@pytest.mark.asyncio +async def test_postgres_sparse_vector_index_requires_hnsw() -> None: + schema = await postgres.TableSchema.from_class(_SparseVectorRow, primary_key=["id"]) + table = postgres.TableTarget(cast(Any, None), schema) + + with pytest.raises(ValueError, match="only supports HNSW"): + table.declare_vector_index(column="embedding", method="ivfflat") diff --git a/python/tests/connectors/test_postgres_target.py b/python/tests/connectors/test_postgres_target.py index 4e23c60cd..f57487ded 100644 --- a/python/tests/connectors/test_postgres_target.py +++ b/python/tests/connectors/test_postgres_target.py @@ -10,7 +10,7 @@ import uuid from dataclasses import dataclass -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, cast if TYPE_CHECKING: import asyncpg @@ -21,7 +21,7 @@ from numpy.typing import NDArray import cocoindex as coco -from cocoindex.resources.schema import VectorSchema +from cocoindex.resources.schema import SparseVector, SparseVectorSchema, VectorSchema from tests import common @@ -194,6 +194,16 @@ class HalfVectorRow: ] +@dataclass +class SparseVectorRow: + id: str + content: str + embedding: Annotated[ + SparseVector, + SparseVectorSchema(size=100), + ] + + @dataclass class TextRow: id: str @@ -223,6 +233,86 @@ def __str__(self) -> str: # ============================================================================= +@pytest.mark.asyncio +async def test_postgres_sparse_vector_round_trip_and_index(pg_env: _PgEnv) -> None: + pool = pg_env.pool + coco_env = pg_env.coco_env + table_name = _unique_name("test_sparsevec") + logical_name = "sparse_ip" + pg_index_name = f"{table_name}__vector__{logical_name}" + source_rows: list[SparseVectorRow] = [ + SparseVectorRow( + id="1", + content="canonical", + embedding=SparseVector(indices=(1, 7), values=(0.5, 0.9)), + ), + SparseVectorRow( + id="2", + content="mapping", + embedding=cast(SparseVector, {4: 0.25, 2: 0.0}), + ), + ] + + try: + + async def declare_fn() -> None: + table = await coco.use_mount( + coco.component_subpath("setup", "table"), + postgres.declare_table_target, + _PG_DB_KEY, + table_name, + await postgres.TableSchema.from_class( + SparseVectorRow, primary_key=["id"] + ), + ) + for row in source_rows: + table.declare_row(row=row) + table.declare_vector_index( + name=logical_name, + column="embedding", + metric="ip", + method="hnsw", + ) + + app = coco.App( + coco.AppConfig(name=f"test_sparsevec_{table_name}", environment=coco_env), + declare_fn, + ) + await app.update() + + async with pool.acquire() as conn: + rows = await conn.fetch( + f'SELECT "id", "embedding"::text AS embedding ' + f'FROM "{table_name}" ORDER BY "id"' + ) + assert [(row["id"], row["embedding"]) for row in rows] == [ + ("1", "{2:0.5,8:0.9}/100"), + ("2", "{5:0.25}/100"), + ] + info = await _index_info(pool, pg_index_name) + assert info is not None and info["amname"] == "hnsw" + assert await _index_opclass_names(pool, pg_index_name) == ["sparsevec_ip_ops"] + + source_rows[:] = [ + SparseVectorRow( + id="1", + content="updated", + embedding=SparseVector(indices=(0, 99), values=(1.0, 0.75)), + ) + ] + await app.update() + + async with pool.acquire() as conn: + rows = await conn.fetch( + f'SELECT "id", "content", "embedding"::text AS embedding ' + f'FROM "{table_name}" ORDER BY "id"' + ) + assert [tuple(row) for row in rows] == [("1", "updated", "{1:1,100:0.75}/100")] + + finally: + await _drop_table(pool, table_name) + + @pytest.mark.asyncio async def test_postgres_declare_vector_index(pg_env: _PgEnv) -> None: """Vector index lifecycle: create with ivfflat → change to hnsw → remove table.""" diff --git a/python/tests/connectors/test_qdrant_target.py b/python/tests/connectors/test_qdrant_target.py index b155d6271..0dc0157eb 100644 --- a/python/tests/connectors/test_qdrant_target.py +++ b/python/tests/connectors/test_qdrant_target.py @@ -47,7 +47,11 @@ _validate_point_id, _vector_params_from_def, ) - from cocoindex.resources.schema import MultiVectorSchema, VectorSchema + from cocoindex.resources.schema import ( + MultiVectorSchema, + SparseVector, + VectorSchema, + ) from tests import common requires_qdrant_url = pytest.mark.skipif( @@ -178,6 +182,19 @@ def test_multivector_schema(self) -> None: @requires_qdrant class TestSparseVectorSupport: + def test_sparse_vector_converter_normalizes_canonical_and_mapping_values( + self, + ) -> None: + canonical = SparseVector(indices=(1, 7), values=(0.5, 0.9)) + + converted = qdrant.sparse_vector(canonical) + assert converted.indices == [1, 7] + assert converted.values == [0.5, 0.9] + + converted_mapping = qdrant.sparse_vector({7: 0.9, 1: 0.5}) + assert converted_mapping.indices == [1, 7] + assert converted_mapping.values == [0.5, 0.9] + @pytest.mark.asyncio async def test_collection_schema_create_resolves_sparse_vector_params(self) -> None: schema = await qdrant.CollectionSchema.create( @@ -426,8 +443,11 @@ async def app_main() -> None: id=1, vector={ "dense": [0.1, 0.2, 0.3, 0.4], - "sparse": qdrant_models.SparseVector( - indices=[1, 7], values=[0.5, 0.9] + "sparse": qdrant.sparse_vector( + SparseVector.from_arrays( + indices=[7, 1], + values=[0.9, 0.5], + ) ), }, payload={"text": "hybrid sparse dense"}, diff --git a/python/tests/connectors/test_snowflake_target.py b/python/tests/connectors/test_snowflake_target.py index ee3d85566..a8b6e6c62 100644 --- a/python/tests/connectors/test_snowflake_target.py +++ b/python/tests/connectors/test_snowflake_target.py @@ -9,7 +9,9 @@ import uuid from typing import Annotated, Any, cast +import numpy as np import pytest +from numpy.typing import NDArray import cocoindex as coco from cocoindex.connectorkits import target @@ -17,6 +19,9 @@ from cocoindex.connectors.snowflake import _target SNOWFLAKE_DB = coco.ContextKey[snowflake.ConnectionConfig]("snowflake_test_db") +_UNPROVIDED_VECTOR_SCHEMA = coco.ContextKey[object]( + "snowflake_test_unprovided_vector_schema" +) @dataclasses.dataclass @@ -43,6 +48,12 @@ class OverrideRow: vector: Annotated[list[float], snowflake.SnowflakeType("ARRAY")] +@dataclasses.dataclass +class ContextAnnotatedVectorRow: + id: int + embedding: Annotated[NDArray[np.float32], _UNPROVIDED_VECTOR_SCHEMA] + + class FakeCursor: def __init__(self) -> None: self.calls: list[tuple[str, tuple[Any, ...] | None]] = [] @@ -126,6 +137,15 @@ async def test_table_target_rejects_invalid_table_identifier() -> None: ) +@pytest.mark.asyncio +async def test_table_schema_does_not_resolve_vector_context_annotations() -> None: + schema = await snowflake.TableSchema.from_class( + ContextAnnotatedVectorRow, primary_key=["id"] + ) + + assert schema.columns["embedding"].type == "VARIANT" + + @pytest.mark.asyncio async def test_table_target_rejects_invalid_column_identifier() -> None: schema = snowflake.TableSchema( diff --git a/python/tests/connectors/test_sqlite_target.py b/python/tests/connectors/test_sqlite_target.py index 6fc864aba..c67425030 100644 --- a/python/tests/connectors/test_sqlite_target.py +++ b/python/tests/connectors/test_sqlite_target.py @@ -29,7 +29,7 @@ # ============================================================================= try: - import sqlite_vec # type: ignore[import-not-found] + __import__("sqlite_vec") HAS_SQLITE_VEC = True except ImportError: diff --git a/python/tests/connectors/test_turbopuffer_target.py b/python/tests/connectors/test_turbopuffer_target.py index 02c0f3595..20f855d8e 100644 --- a/python/tests/connectors/test_turbopuffer_target.py +++ b/python/tests/connectors/test_turbopuffer_target.py @@ -20,7 +20,7 @@ import pytest_asyncio import cocoindex as coco -from cocoindex.resources.schema import VectorSchema +from cocoindex.resources.schema import SparseVector, SparseVectorSchema, VectorSchema from tests import common @@ -134,6 +134,14 @@ def test_single_vector_dict_rejected(self, single_schema: Any) -> None: with pytest.raises(ValueError, match="single unnamed vector"): _row_to_upsert(row, single_schema) + def test_sparse_vector_rejected(self, single_schema: Any) -> None: + row = turbopuffer.Row( + id="a", + vector=SparseVector(indices=(1,), values=(0.5,)), # type: ignore[arg-type] + ) + with pytest.raises(ValueError, match="does not support sparse vector"): + _row_to_upsert(row, single_schema) + def test_named_vectors(self, named_schema: Any) -> None: row = turbopuffer.Row( id="a", @@ -148,6 +156,17 @@ def test_named_vectors(self, named_schema: Any) -> None: "title": "T", } + def test_named_sparse_mapping_rejected(self, named_schema: Any) -> None: + row = turbopuffer.Row( + id="a", + vector={ + "text": {1: 0.5, 7: 0.9}, # type: ignore[dict-item] + "image": [0.5, 0.5], + }, + ) + with pytest.raises(ValueError, match="does not support sparse vector"): + _row_to_upsert(row, named_schema) + def test_named_vectors_missing_field(self, named_schema: Any) -> None: row = turbopuffer.Row(id="a", vector={"text": [1.0, 2.0, 3.0]}) with pytest.raises(ValueError, match="missing vector fields"): @@ -186,6 +205,15 @@ def test_named_vector_attribute_collision(self, named_schema: Any) -> None: @requires_turbopuffer class TestNamespaceSchemaCreate: + @pytest.mark.asyncio + async def test_sparse_schema_rejected(self) -> None: + with pytest.raises(ValueError, match="does not support sparse vector"): + await turbopuffer.NamespaceSchema.create( + vectors=turbopuffer.VectorDef( + schema=SparseVectorSchema(size=100) # type: ignore[arg-type] + ) + ) + @pytest.mark.asyncio async def test_named_vector_id_rejected(self) -> None: # "id" as a named vector field name would silently overwrite row.id diff --git a/python/tests/connectors/test_valkey_target.py b/python/tests/connectors/test_valkey_target.py index 6a0209998..e5ca6e286 100644 --- a/python/tests/connectors/test_valkey_target.py +++ b/python/tests/connectors/test_valkey_target.py @@ -14,14 +14,14 @@ import struct import uuid from collections.abc import AsyncIterator, Iterator -from typing import Any, NamedTuple +from typing import Any, NamedTuple, cast import numpy as np import pytest import pytest_asyncio import cocoindex as coco -from cocoindex.resources.schema import VectorSchema +from cocoindex.resources.schema import SparseVector, VectorSchema from tests import common @@ -39,6 +39,7 @@ if HAS_GLIDE: from cocoindex.connectors import valkey + from cocoindex.connectors.valkey._target import _vector_to_bytes requires_glide = pytest.mark.skipif(not HAS_GLIDE, reason="valkey-glide not installed") @@ -61,6 +62,21 @@ _VECTOR_SCHEMA = VectorSchema(dtype=np.dtype(np.float32), size=_DIM) +@requires_glide +@pytest.mark.parametrize( + "vector", + [ + {1: 0.5, 7: 0.9}, + SparseVector(indices=(1, 7), values=(0.5, 0.9)), + ], +) +def test_vector_to_bytes_rejects_sparse_shapes(vector: object) -> None: + with pytest.raises( + ValueError, match="Valkey does not support sparse vector values" + ): + _vector_to_bytes(cast(Any, vector)) + + def _unique_name(prefix: str) -> str: return f"{prefix}_{uuid.uuid4().hex[:8]}" diff --git a/python/tests/connectors/test_zvec_target.py b/python/tests/connectors/test_zvec_target.py index a7c17d713..4a896cc56 100644 --- a/python/tests/connectors/test_zvec_target.py +++ b/python/tests/connectors/test_zvec_target.py @@ -8,7 +8,7 @@ import uuid from dataclasses import dataclass from pathlib import Path -from typing import Annotated, Any, Iterator +from typing import Annotated, Any, Iterator, cast import numpy as np import pytest @@ -17,7 +17,7 @@ import cocoindex as coco from cocoindex._internal.context_keys import ContextProvider from cocoindex.connectorkits import target -from cocoindex.resources.schema import VectorSchema +from cocoindex.resources.schema import SparseVector, SparseVectorSchema, VectorSchema from tests import common @@ -25,6 +25,7 @@ import zvec from cocoindex.connectors import zvec as zc + from cocoindex.connectors.zvec._target import _to_float_list, _to_sparse_dict HAS_ZVEC = True except ImportError: @@ -120,6 +121,22 @@ class SparseDoc: sparse: Annotated[dict[int, float], zc.ZvecVectorDef(sparse=True)] +@dataclass +class CanonicalSparseDoc: + id: str + title: str + sparse: SparseVector + + +@dataclass +class AnnotatedSparseDoc: + id: str + sparse: Annotated[ + SparseVector, + SparseVectorSchema(size=10_000), + ] + + @dataclass class MultiVectorDoc: id: str @@ -140,6 +157,12 @@ class Fp16Doc: embedding: _Embedding16 +@dataclass +class ListVectorDoc: + id: str + embedding: Annotated[list[float], VectorSchema(dtype=np.dtype(np.float32), size=4)] + + @dataclass class QuantizedDoc: id: str @@ -365,6 +388,21 @@ def test_dense_vector(conn: Any) -> None: assert [d.id for d in results] == ["1"] +def test_list_dense_vector_compatibility(conn: Any) -> None: + _reset(ListVectorDoc, "test_list_dense") + app = _make_app(conn, "test_list_dense_vector_compatibility") + + _rows.append(ListVectorDoc(id="1", embedding=[0.1, 0.2, 0.3, 0.4])) + app.update_blocking() + + col = conn.open_existing("test_list_dense") + results = col.query( + zvec.Query(field_name="embedding", vector=[0.1, 0.2, 0.3, 0.4]), + topk=5, + ) + assert [d.id for d in results] == ["1"] + + def test_sparse_vector(conn: Any) -> None: _reset(SparseDoc, "test_sparse") app = _make_app(conn, "test_sparse_vector") @@ -383,6 +421,56 @@ def test_sparse_vector(conn: Any) -> None: assert [d.id for d in results] == ["1"] +def test_canonical_sparse_vector_and_mapping_values(conn: Any) -> None: + _reset(CanonicalSparseDoc, "test_canonical_sparse") + app = _make_app(conn, "test_canonical_sparse_vector") + + _rows.extend( + [ + CanonicalSparseDoc( + id="1", + title="canonical", + sparse=SparseVector(indices=(1, 7), values=(0.5, 0.9)), + ), + CanonicalSparseDoc( + id="2", + title="mapping", + sparse=cast(SparseVector, {9: 0.8, 2: 0.4}), + ), + ] + ) + app.update_blocking() + + col = conn.open_existing("test_canonical_sparse") + canonical_results = col.query( + zvec.Query(field_name="sparse", vector={1: 0.5, 7: 0.9}), topk=5 + ) + mapping_results = col.query( + zvec.Query(field_name="sparse", vector={2: 0.4, 9: 0.8}), topk=5 + ) + assert canonical_results[0].id == "1" + assert mapping_results[0].id == "2" + + +def test_sparse_mapping_is_normalized_before_zvec_write() -> None: + assert list(_to_sparse_dict({7: 0.9, 1: 0.5}).items()) == [ + (1, 0.5), + (7, 0.9), + ] + + +@pytest.mark.parametrize( + "value", + [ + {1: 0.5, 7: 0.9}, + SparseVector(indices=(1, 7), values=(0.5, 0.9)), + ], +) +def test_dense_value_conversion_rejects_sparse_shapes(value: object) -> None: + with pytest.raises(ValueError, match="does not support sparse values"): + _to_float_list(value) + + def test_fts_field(conn: Any) -> None: _reset(FtsDoc, "test_fts") app = _make_app(conn, "test_fts_field") @@ -567,6 +655,58 @@ class Float64Doc: with pytest.raises(ValueError, match="float32 or float16"): await zc.CollectionSchema.from_class(Float64Doc, primary_key=["id"]) + canonical_schema = await zc.CollectionSchema.from_class( + CanonicalSparseDoc, primary_key=["id"] + ) + annotated_schema = await zc.CollectionSchema.from_class( + AnnotatedSparseDoc, primary_key=["id"] + ) + override_schema = await zc.CollectionSchema.from_class( + CanonicalSparseDoc, + primary_key=["id"], + column_overrides={"sparse": SparseVectorSchema(size=None)}, + ) + for sparse_schema in (canonical_schema, annotated_schema, override_schema): + assert sparse_schema.columns["sparse"].kind == "sparse" + assert ( + sparse_schema.columns["sparse"].data_type + == zvec.DataType.SPARSE_VECTOR_FP32 + ) + + @dataclass + class NullableSparseDoc: + id: str + sparse: Annotated[SparseVector | None, SparseVectorSchema()] + + nullable_schema = await zc.CollectionSchema.from_class( + NullableSparseDoc, primary_key=["id"] + ) + assert nullable_schema.columns["sparse"].kind == "sparse" + assert nullable_schema.columns["sparse"].nullable + + @dataclass + class InvalidLegacySparseDoc: + id: str + sparse: Annotated[str, zc.ZvecVectorDef(sparse=True)] + + with pytest.raises(ValueError, match="requires a SparseVector or mapping field"): + await zc.CollectionSchema.from_class(InvalidLegacySparseDoc, primary_key=["id"]) + + @dataclass + class NativeAndSparseDoc: + id: str + sparse: Annotated[ + SparseVector, + SparseVectorSchema(), + ] + + with pytest.raises(ValueError, match="cannot combine ZvecType"): + await zc.CollectionSchema.from_class( + NativeAndSparseDoc, + primary_key=["id"], + column_overrides={"sparse": zc.ZvecType(zvec.DataType.STRING)}, + ) + def test_fp16_dense_vector(conn: Any) -> None: _reset(Fp16Doc, "test_fp16") diff --git a/python/tests/core/test_datatype.py b/python/tests/core/test_datatype.py index 0cdeb7839..9c4940449 100644 --- a/python/tests/core/test_datatype.py +++ b/python/tests/core/test_datatype.py @@ -274,3 +274,14 @@ def test_type_with_attributes() -> None: annotations=("Annotation1", "Annotation2"), nullable=False, ) + + +def test_nullable_type_preserves_annotations() -> None: + typ = Annotated[Annotated[str, "Inner"] | None, "Outer"] + + result = analyze_type_info(typ) + + assert result.core_type is str + assert result.base_type is str + assert result.nullable + assert result.annotations == ("Inner", "Outer") diff --git a/python/tests/resources/test_sparse_vector.py b/python/tests/resources/test_sparse_vector.py new file mode 100644 index 000000000..106118ec0 --- /dev/null +++ b/python/tests/resources/test_sparse_vector.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import numpy as np +import pytest + +import cocoindex as coco +from cocoindex.resources.schema import ( + SparseVector, + SparseVectorSchema, + as_sparse_vector, + get_sparse_vector_schema, +) + +from tests import common + + +class _SparseSchemaProvider: + def __init__(self, schema: SparseVectorSchema) -> None: + self._schema = schema + + async def __coco_sparse_vector_schema__(self) -> SparseVectorSchema: + return self._schema + + +_SPARSE_SCHEMA_KEY = coco.ContextKey[_SparseSchemaProvider]( + "test_sparse_vector/schema_provider" +) + + +def test_sparse_vector_from_arrays_normalizes_and_sorts() -> None: + sparse_vector = SparseVector.from_arrays( + np.array([7, 1], dtype=np.int64), + np.array([0.5, 1], dtype=np.float32), + ) + + assert sparse_vector == SparseVector(indices=(1, 7), values=(1.0, 0.5)) + + +def test_sparse_vector_from_mapping_normalizes_order_and_values() -> None: + sparse_vector = SparseVector.from_mapping({7: 0.9, 1: 0.5}) + + assert sparse_vector == SparseVector(indices=(1, 7), values=(0.5, 0.9)) + assert as_sparse_vector(sparse_vector) is sparse_vector + + +def test_sparse_vector_from_mapping_requires_integral_keys() -> None: + with pytest.raises(TypeError, match="indices must be integers"): + SparseVector.from_mapping({1.5: 0.5}) # type: ignore[dict-item] + + with pytest.raises(TypeError, match="not bool"): + SparseVector.from_mapping({True: 0.5}) + + for value in (True, "0.5"): + with pytest.raises(TypeError, match="values must be real numbers"): + SparseVector.from_mapping({1: value}) # type: ignore[dict-item] + + +@pytest.mark.parametrize( + ("indices", "values", "message"), + [ + ((1,), (), "same length"), + ((2, 1), (0.2, 0.1), "sorted ascending and unique"), + ((1, 1), (0.1, 0.2), "sorted ascending and unique"), + ((-1,), (0.1,), "non-negative"), + ], +) +def test_sparse_vector_rejects_invalid_indices_and_values( + indices: tuple[int, ...], values: tuple[float, ...], message: str +) -> None: + with pytest.raises(ValueError, match=message): + SparseVector(indices=indices, values=values) + + +@pytest.mark.parametrize( + ("indices", "values", "message"), + [ + ([1], (0.5,), "indices must be a tuple"), + (np.array([1]), (0.5,), "indices must be a tuple"), + ((1,), [0.5], "values must be a tuple"), + ((1,), np.array([0.5]), "values must be a tuple"), + ((1.0,), (0.5,), "indices must contain only integers"), + ((1,), (1,), "values must contain only floats"), + ((True,), (0.5,), "indices must contain only integers"), + ], +) +def test_sparse_vector_rejects_invalid_containers_and_elements( + indices: object, values: object, message: str +) -> None: + with pytest.raises(TypeError, match=message): + SparseVector(indices=indices, values=values) # type: ignore[arg-type] + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_sparse_vector_rejects_non_finite_values(value: float) -> None: + with pytest.raises(ValueError, match="values must be finite"): + SparseVector(indices=(1,), values=(value,)) + with pytest.raises(ValueError, match="values must be finite"): + SparseVector.from_mapping({1: value}) + + +@pytest.mark.parametrize("value", [[(1, 0.5)], np.array([0.5]), "1:0.5"]) +def test_as_sparse_vector_rejects_non_mapping_values(value: object) -> None: + with pytest.raises(TypeError, match="SparseVector or Mapping"): + as_sparse_vector(value) # type: ignore[arg-type] + + +def test_sparse_vector_mapping_normalization_is_canonical() -> None: + first = {1: 0.5, 7: 0.9} + second = {7: 0.9, 1: 0.5} + + assert as_sparse_vector(first) == as_sparse_vector(second) + + +def test_sparse_vector_schema_defaults_and_validates_size() -> None: + assert SparseVectorSchema(size=100).size == 100 + numpy_size_schema = SparseVectorSchema(size=np.int64(100)) # type: ignore[arg-type] + assert numpy_size_schema.size == 100 + assert type(numpy_size_schema.size) is int + + for size in (0, -1): + with pytest.raises(ValueError, match="size must be positive"): + SparseVectorSchema(size=size) + + with pytest.raises(TypeError, match="size must be an integer"): + SparseVectorSchema(size=1.5) # type: ignore[arg-type] + + +@pytest.mark.asyncio +async def test_get_sparse_vector_schema_resolves_schema_provider_and_context_key() -> ( + None +): + schema = SparseVectorSchema(size=100) + provider = _SparseSchemaProvider(schema) + + assert await get_sparse_vector_schema(schema) is schema + assert await get_sparse_vector_schema(provider) is schema + + env = common.create_test_env(__file__) + env.context_provider.provide(_SPARSE_SCHEMA_KEY, provider) + + async def resolve_from_context() -> SparseVectorSchema | None: + return await get_sparse_vector_schema(_SPARSE_SCHEMA_KEY) + + app = coco.App( + coco.AppConfig(name="test_sparse_schema_context", environment=env), + resolve_from_context, + ) + assert await app.update() is schema