From 411a526efa439276cb3f99ec25607562197b5e5f Mon Sep 17 00:00:00 2001 From: Daniele Date: Sun, 15 Mar 2026 10:36:52 +0100 Subject: [PATCH 1/8] Add pgvector vectorstore to workspace Include datapizza-ai-vectorstores-pgvector in the workspace: added the package to the members list and declared datapizza-ai-vectorstores-pgvector = {workspace = true} so the pgvector vectorstore is part of the monorepo build and dependency graph. --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 14037a35..91b05f77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,7 @@ members = [ # Vectorstores "datapizza-ai-vectorstores/datapizza-ai-vectorstores-qdrant", "datapizza-ai-vectorstores/datapizza-ai-vectorstores-milvus", + "datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector", # Cache "datapizza-ai-cache/redis", @@ -102,6 +103,7 @@ datapizza-ai-rerankers-together = {workspace = true} # Vectorstores datapizza-ai-vectorstores-qdrant = {workspace = true} +datapizza-ai-vectorstores-pgvector = {workspace = true} # Cache datapizza-ai-cache-redis = {workspace = true} From bddfa8b76acbc38cedc64597630f2d57f350ef63 Mon Sep 17 00:00:00 2001 From: Daniele Date: Sun, 15 Mar 2026 10:37:45 +0100 Subject: [PATCH 2/8] Add pgvector vectorstore documentation Add new docs/API Reference/Vectorstore/pgvector_vectorstore.md documenting the PgVectorVectorstore. Includes installation, basic usage examples, optional psycopg_pool connection pooling, JSONB metadata filtering, batch insert behavior (executemany), clean shutdown (close/aclose), and collection/index management examples. --- .../Vectorstore/pgvector_vectorstore.md | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 docs/API Reference/Vectorstore/pgvector_vectorstore.md diff --git a/docs/API Reference/Vectorstore/pgvector_vectorstore.md b/docs/API Reference/Vectorstore/pgvector_vectorstore.md new file mode 100644 index 00000000..6ced3760 --- /dev/null +++ b/docs/API Reference/Vectorstore/pgvector_vectorstore.md @@ -0,0 +1,136 @@ +# pgvector + +```python +pip install datapizza-ai-vectorstores-pgvector +``` + + +::: datapizza.vectorstores.pgvector.PgVectorVectorstore + options: + show_source: false + +## Usage + +```python +from datapizza.core.vectorstore import VectorConfig +from datapizza.vectorstores.pgvector import PgVectorVectorstore + +store = PgVectorVectorstore(dsn="postgresql://user:pass@localhost:5432/db") + +store.create_collection( + "docs", + vector_config=[VectorConfig(name="embedding", dimensions=1536)], +) + +# Optional: create an ANN index immediately +store.create_collection( + "docs_with_index", + vector_config=[VectorConfig(name="embedding", dimensions=1536)], + create_index=True, +) + +# Query the collection +results = store.search( + collection_name="docs", + query_vector=[0.1, 0.2, ...], + k=5, +) +``` + +## Advanced Usage + +### Optional connection pooling + +The vectorstore supports optional connection pooling using `psycopg_pool`. This is useful in multi-threaded or high-concurrency environments. + +```python +store = PgVectorVectorstore( + dsn="postgresql://user:pass@localhost:5432/db", + use_pool=True, + pool_min_size=2, + pool_max_size=10, +) +``` + +> Note: `psycopg_pool` is optional. If `use_pool=True` and it is not installed, the constructor will raise an `ImportError`. + +### JSONB filtering (metadata) + +You can filter operations on the `metadata` column using JSONB containment. This works for `search`, `retrieve`, `update`, and `remove`. + +```python +# Search only items with metadata.tag == "keep" +results = store.search( + collection_name="docs", + query_vector=[...], + k=10, + filters={"tag": "keep"}, +) + +# Update only filtered items +store.update( + "docs", + payload={"status": "updated"}, + points=[...], + filters={"tag": "keep"}, +) + +# Remove only filtered items +store.remove( + "docs", + ids=[...], + filters={"tag": "skip"}, +) +``` + +## Performance & cleanup + +### Batch inserts (executemany) + +`add()` accepts either a single `Chunk` or a list of `Chunk` objects. When given a list, the implementation uses `executemany` for a batched insert, which is significantly faster for large uploads. + +```python +chunks = [ + Chunk( + id=str(uuid.uuid4()), + text=f"batch {i}", + embeddings=[DenseEmbedding(name="embedding", vector=[float(i), float(i)])], + ) + for i in range(100) +] + +store.add(chunks, collection_name="docs") +``` + +### Clean shutdown (pool close) + +If you use connection pooling, it is recommended to close resources when your application shuts down. + +```python +store.close() + +# Async +await store.aclose() +``` + +## Collection management + +```python +store.list_collections() +store.describe_collection("docs") +store.drop_collection("docs") +``` + +## Index management + +```python +store.create_index( + "docs", + method="hnsw", + index_name="docs_embedding_hnsw_idx", +) + +store.list_indexes("docs") +store.describe_index("docs_embedding_hnsw_idx") +store.drop_index("docs_embedding_hnsw_idx") +``` From 7f9a918d58127078c02c7b20207c039f53c9f3d4 Mon Sep 17 00:00:00 2001 From: Daniele Date: Sun, 15 Mar 2026 10:38:49 +0100 Subject: [PATCH 3/8] Add pyproject and pytest config for pgvector store Add a new package config for datapizza-ai-vectorstores-pgvector --- .../pyproject.toml | 61 +++++++++++++++++++ .../pytest.ini | 5 ++ 2 files changed, 66 insertions(+) create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pyproject.toml create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pytest.ini diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pyproject.toml b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pyproject.toml new file mode 100644 index 00000000..fe5efb96 --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pyproject.toml @@ -0,0 +1,61 @@ +# Build system configuration +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +# Project metadata +[project] +name = "datapizza-ai-vectorstores-pgvector" +version = "0.0.1" +description = "pgvector vectorstore for the datapizza-ai framework" +readme = "README.md" +license = {text = "MIT"} + +requires-python = ">=3.10.0,<4" +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Application Frameworks", +] +dependencies = [ + "datapizza-ai-core>=0.1.0,<0.2.0", + "psycopg[binary]>=3.1", + "pgvector>=0.3", +] + +# Development dependencies +[dependency-groups] +dev = [ + "deptry>=0.23.0", + "pytest", + "pytest-asyncio", + "testcontainers", + "ruff>=0.11.5", +] + +# Hatch build configuration +[tool.hatch.build.targets.sdist] +include = ["datapizza"] +exclude = ["**/BUILD"] + +[tool.hatch.build.targets.wheel] +include = ["datapizza"] +exclude = ["**/BUILD"] + +# Ruff configuration +[tool.ruff] +line-length = 88 + +[tool.ruff.lint] +select = [ + # "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "B", # flake8-bugbear + "I", # isort + "UP", # pyupgrade + "SIM", # flake8-simplify + "RUF", # Ruff-specific rules + "C4", # flake8-comprehensions +] diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pytest.ini b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pytest.ini new file mode 100644 index 00000000..d7a3783e --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +markers = + integration: mark a test as an integration test that may require external resources. + asyncio: mark tests that use pytest-asyncio + unit: mark unit tests From 196acb3b8e0aa792714d4e4dba80e68c4eabea69 Mon Sep 17 00:00:00 2001 From: Daniele Date: Sun, 15 Mar 2026 10:39:18 +0100 Subject: [PATCH 4/8] Add README for pgvector vectorstore Introduce a new README for the datapizza-ai-pgvector vectorstore implementation. --- .../README.md | 117 ++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/README.md diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/README.md b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/README.md new file mode 100644 index 00000000..b9a7991d --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/README.md @@ -0,0 +1,117 @@ +# datapizza-ai-vectorstores-pgvector + +PostgreSQL + pgvector vectorstore implementation for the datapizza-ai framework. + +This package provides a `PgVectorVectorstore` that implements the `Vectorstore` interface from `datapizza-ai-core`. + +## Quickstart + +```python +from datapizza.core.vectorstore import VectorConfig +from datapizza.vectorstores.pgvector import PgVectorVectorstore + +store = PgVectorVectorstore(dsn="postgresql://user:pass@localhost:5432/db") + +# Create a collection (table) with a vector column +store.create_collection("docs", vector_config=[VectorConfig(name="embedding", dimensions=1536)]) + +# Optional: create an ANN index immediately (HNSW by default) +store.create_collection( + "docs_with_index", + vector_config=[VectorConfig(name="embedding", dimensions=1536)], + create_index=True, +) + +# Query the collection +results = store.search( + collection_name="docs", + query_vector=[0.1, 0.2, ...], + k=5, +) + +# Drop the collection (deletes the table) +store.drop_collection("docs") + +# Inspect existing collections +print(store.list_collections()) +print(store.describe_collection("docs")) +``` + +## Advanced Usage + +### Optional connection pooling (enterprise) + +`PgVectorVectorstore` supports optional connection pooling using `psycopg_pool`. + +```python +from datapizza.vectorstores.pgvector import PgVectorVectorstore + +store = PgVectorVectorstore( + dsn="postgresql://user:pass@localhost:5432/db", + use_pool=True, + pool_min_size=2, + pool_max_size=10, +) +``` + +> Note: `psycopg_pool` is optional. If `use_pool=True` and it is not installed, the constructor will raise an `ImportError`. + +### Metadata filtering (JSONB) + +You can filter search/update/remove/retrieve operations using JSONB filtering on the `metadata` column. + +```python +# Only search in items where metadata contains {"tag": "keep"} +results = store.search( + collection_name="docs", + query_vector=[...], + k=10, + filters={"tag": "keep"}, +) + +# Update only items matching a filter +store.update( + "docs", + payload={"status": "updated"}, + points=[...], + filters={"tag": "keep"}, +) + +# Remove items matching a filter +store.remove( + "docs", + ids=[...], + filters={"tag": "skip"}, +) +``` + +## Performance & cleanup + +### Batch inserts (executemany) + +`add()` accepts a single `Chunk` or a list of `Chunk` objects. When passing a list, the store performs a batched insert using `executemany`, which is much faster for large uploads. + +```python +chunks = [ + Chunk( + id=str(uuid.uuid4()), + text=f"batch {i}", + embeddings=[DenseEmbedding(name="embedding", vector=[float(i), float(i)])], + ) + for i in range(100) +] + +store.add(chunks, collection_name="docs") +``` + +### Clean shutdown (pool close) + +If you enable pooling, it's a good idea to explicitly close resources when your application shuts down. + +```python +store.close() + +# Async +await store.aclose() +``` +``` From a3f2473cb1bb94bf537b6e01a6028c2977b7744c Mon Sep 17 00:00:00 2001 From: Daniele Date: Sun, 15 Mar 2026 10:41:26 +0100 Subject: [PATCH 5/8] Add PgVector vectorstore implementation Introduce PgVectorVectorstore: a Postgres + pgvector-backed implementation of the Vectorstore interface. Provides sync/async APIs and optional connection pooling, DSN normalization, safe schema/table quoting, and local + DB-backed collection config caching. --- .../pgvector/pgvector_vectorstore.py | 1003 +++++++++++++++++ 1 file changed, 1003 insertions(+) create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py new file mode 100644 index 00000000..4862341d --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py @@ -0,0 +1,1003 @@ +"""PostgreSQL + pgvector vectorstore implementation. + +This implementation is intended to match the `Vectorstore` interface from +`datapizza-ai-core`, allowing users to store and retrieve embeddings using +Postgres + pgvector. +""" + +from __future__ import annotations + +import contextlib +import logging +import re +from typing import Any, cast, Dict, List, Optional + +import psycopg +from psycopg.rows import dict_row +from psycopg.types.json import Jsonb +from pgvector.psycopg import Vector, register_vector, register_vector_async + +try: + from psycopg_pool import AsyncConnectionPool, ConnectionPool +except ImportError: # pragma: no cover + ConnectionPool = None # type: ignore[assignment] + AsyncConnectionPool = None # type: ignore[assignment] + +from datapizza.core.vectorstore import Distance, VectorConfig, Vectorstore +from datapizza.type import Chunk, DenseEmbedding, EmbeddingFormat, SparseEmbedding + +log = logging.getLogger(__name__) + +_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _validate_identifier(name: str) -> str: + if not _IDENTIFIER_RE.match(name): + raise ValueError( + f"Invalid identifier '{name}'. Only alphanumeric characters and underscores are allowed, and it must not start with a digit." + ) + return name + + +class PgVectorVectorstore(Vectorstore): + """A vectorstore implementation backed by Postgres + pgvector.""" + + def __init__( + self, + dsn: str | None = None, + host: str | None = None, + port: int | None = None, + database: str | None = None, + user: str | None = None, + password: str | None = None, + schema: str = "public", + use_pool: bool = False, + pool_min_size: int = 1, + pool_max_size: int = 10, + pool_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + """Initialize the PgVectorVectorstore.""" + self.dsn = dsn + self.host = host + self.port = port + self.database = database + self.user = user + self.password = password + self.schema = schema + self.use_pool = use_pool + self.pool_min_size = pool_min_size + self.pool_max_size = pool_max_size + self.pool_kwargs = pool_kwargs or {} + self.kwargs: dict[str, Any] = kwargs + + # Cached connections (if not using a pool) + self._conn: Optional[psycopg.Connection] = None + self._a_conn: Optional[psycopg.AsyncConnection] = None + + # Optional pools + self._pool: Optional[ConnectionPool] = None + self._a_pool: Optional[AsyncConnectionPool] = None + + # Local cache of collection vector configs (for name inference) + self._collections: Dict[str, List[VectorConfig]] = {} + + def _qualified_table(self, collection_name: str) -> str: + """Return a safely quoted schema-qualified table name.""" + schema = _validate_identifier(self.schema) + """Return a safely quoted schema-qualified table name.""" + schema = _validate_identifier(self.schema) + table = _validate_identifier(collection_name) + return f'"{schema}"."{table}"' + + def _execute(self, cur: Any, query: Any, params: Any = None): + """Execute on a cursor with relaxed typing for psycopg's execute signature.""" + if params is None: + return cast(Any, cur).execute(query) + return cast(Any, cur).execute(query, params) + + async def _a_execute(self, cur: Any, query: Any, params: Any = None): + """Async execute wrapper with relaxed typing.""" + if params is None: + return await cast(Any, cur).execute(query) + return await cast(Any, cur).execute(query, params) + + def _normalize_dsn(self, dsn: str) -> str: + """Normalize DSNs coming from SQLAlchemy/testcontainers.""" + if dsn.startswith("postgresql+psycopg2://"): + return dsn.replace("postgresql+psycopg2://", "postgresql://", 1) + return dsn + + def _build_dsn(self) -> str: + """Construct a libpq-style DSN string from provided connection components.""" + if self.dsn: + return self._normalize_dsn(self.dsn) + + parts: list[str] = [] + if self.host: + parts.append(f"host={self.host}") + if self.port: + parts.append(f"port={self.port}") + if self.database: + parts.append(f"dbname={self.database}") + if self.user: + parts.append(f"user={self.user}") + if self.password: + parts.append(f"password={self.password}") + return " ".join(parts) + + def _get_pool(self) -> ConnectionPool: + if not self.use_pool: + raise RuntimeError("Connection pooling is disabled") + if ConnectionPool is None: + raise ImportError( + "psycopg_pool is required for connection pooling (pip install psycopg_pool)" + ) + if self._pool is None: + dsn = self._build_dsn() + self._pool = ConnectionPool( + dsn, + min_size=self.pool_min_size, + max_size=self.pool_max_size, + open=True, + **self.pool_kwargs, + ) + return self._pool + + def _conn_context(self): + """Return a context manager yielding a connection (pooled or single).""" + if self.use_pool: + pool = self._get_pool() + + @contextlib.contextmanager + def _ctx(): + with pool.connection() as conn: + # Ensure pgvector is registered on each pooled connection. + register_vector(conn) + yield conn + + return _ctx() + + @contextlib.contextmanager + def _ctx(): + conn = self._get_conn() + yield conn + + return _ctx() + + async def _get_a_pool(self) -> AsyncConnectionPool: + if not self.use_pool: + raise RuntimeError("Connection pooling is disabled") + if AsyncConnectionPool is None: + raise ImportError( + "psycopg_pool is required for async connection pooling (pip install psycopg_pool)" + ) + if self._a_pool is None: + dsn = self._build_dsn() + # Avoid deprecated behavior: do not open the pool in the constructor. + self._a_pool = AsyncConnectionPool( + dsn, + min_size=self.pool_min_size, + max_size=self.pool_max_size, + open=False, + **self.pool_kwargs, + ) + await self._a_pool.open() + return self._a_pool + + async def _a_conn_context(self): + """Return an async context manager yielding a connection.""" + if self.use_pool: + pool = await self._get_a_pool() + + @contextlib.asynccontextmanager + async def _ctx(): + async with pool.connection() as conn: + # Ensure pgvector is registered on each pooled connection. + await register_vector_async(conn) + yield conn + + return _ctx() + + @contextlib.asynccontextmanager + async def _ctx(): + conn = await self._get_a_conn() + yield conn + + return _ctx() + + def _build_filter_clause( + self, filters: Optional[Dict[str, Any]] = None, prefix: str = "WHERE" + ) -> tuple[str, list[Any]]: + """Build an optional SQL clause for JSONB metadata filters. + + Args: + filters: A dict of metadata fields to match via JSONB containment. + prefix: SQL clause prefix, e.g. "WHERE" or "AND". + """ + if not filters: + return "", [] + + # Use JSONB containment for simple metadata filtering + return f"{prefix} metadata @> %s", [Jsonb(filters)] + + def _get_conn(self) -> psycopg.Connection: + if self._conn is None: + conn_args = self.kwargs.copy() + + if self.dsn: + # psycopg expects a normal libpq connection string; testcontainers gives + # SQLAlchemy-style strings like postgresql+psycopg2://... + dsn = self.dsn + if dsn.startswith("postgresql+psycopg2://"): + dsn = dsn.replace("postgresql+psycopg2://", "postgresql://", 1) + conn = psycopg.connect(dsn, **conn_args) + else: + if self.host: + conn_args["host"] = self.host + if self.port: + conn_args["port"] = self.port + if self.database: + conn_args["dbname"] = self.database + if self.user: + conn_args["user"] = self.user + if self.password: + conn_args["password"] = self.password + conn = psycopg.connect(**conn_args) + + register_vector(conn) + conn.autocommit = True + self._conn = conn + return self._conn + + async def _get_a_conn(self) -> psycopg.AsyncConnection: + if self._a_conn is None: + conn_args = self.kwargs.copy() + + if self.dsn: + dsn = self.dsn + if dsn.startswith("postgresql+psycopg2://"): + dsn = dsn.replace("postgresql+psycopg2://", "postgresql://", 1) + conn = await psycopg.AsyncConnection.connect(dsn, **conn_args) + else: + if self.host: + conn_args["host"] = self.host + if self.port: + conn_args["port"] = self.port + if self.database: + conn_args["dbname"] = self.database + if self.user: + conn_args["user"] = self.user + if self.password: + conn_args["password"] = self.password + conn = await psycopg.AsyncConnection.connect(**conn_args) + + # Use the async registration helper for async connections + await register_vector_async(conn) + await conn.set_autocommit(True) + self._a_conn = conn + return self._a_conn + + def create_collection( + self, + collection_name: str, + vector_config: list[VectorConfig], + create_index: bool = False, + index_method: str = "hnsw", + index_name: Optional[str] = None, + ivfflat_lists: Optional[int] = None, + **kwargs, + ): + """Create a table for the given collection name and vector configuration. + + Args: + collection_name: Name of the collection/table. + vector_config: List of VectorConfig describing columns. + create_index: If True, create an ANN index immediately. + index_method: One of "hnsw" or "ivfflat". + index_name: Optional explicit index name. + ivfflat_lists: Required when using ivfflat. + """ + if not collection_name: + raise ValueError("collection_name must be provided") + + table = self._qualified_table(collection_name) + + # Persist config locally for name inference + self._collections[collection_name] = vector_config + + # Persist config to the metadata table so we can rebuild the cache later + meta = [ + { + "name": cfg.name, + "dimensions": cfg.dimensions, + "distance": cfg.distance.value if hasattr(cfg, "distance") else None, + "format": cfg.format.value if hasattr(cfg, "format") else None, + } + for cfg in vector_config + ] + self._write_meta_config(collection_name, meta) + + # Build schema + cols = [ + "id TEXT PRIMARY KEY", + "text TEXT NOT NULL", + "metadata JSONB NOT NULL DEFAULT '{}'", + ] + + for cfg in vector_config: + if cfg.format != EmbeddingFormat.DENSE: + log.warning( + "Skipping non-dense vector config %s: pgvector only supports dense embeddings.", + cfg, + ) + continue + name = cfg.name or "embedding" + _validate_identifier(name) + cols.append(f'"{name}" vector({cfg.dimensions})') + + with self._conn_context() as conn: + with conn.cursor() as cur: + self._execute(cur, f"CREATE SCHEMA IF NOT EXISTS \"{self.schema}\";") + self._execute(cur, f"CREATE TABLE IF NOT EXISTS {table} ({', '.join(cols)});") + + if create_index: + self.create_index( + collection_name=collection_name, + vector_name=None, + method=index_method, + index_name=index_name, + ivfflat_lists=ivfflat_lists, + ) + + def _chunk_to_row(self, chunk: Chunk) -> Dict[str, Any]: + if not chunk.embeddings: + raise ValueError("Chunk must have an embedding") + + row: Dict[str, Any] = { + "id": chunk.id, + "text": chunk.text or "", + "metadata": Jsonb(chunk.metadata or {}), + } + + for emb in chunk.embeddings: + if isinstance(emb, DenseEmbedding): + name = emb.name or "embedding" + row[name] = Vector(emb.vector) + elif isinstance(emb, SparseEmbedding): + raise ValueError("Sparse embeddings are not supported by pgvector") + else: + raise ValueError(f"Unsupported embedding type: {type(emb)}") + + return row + + def add(self, chunk: Chunk | list[Chunk], collection_name: str | None = None): + if not collection_name: + raise ValueError("collection_name must be provided") + + chunks = [chunk] if isinstance(chunk, Chunk) else chunk + if not chunks: + return + + rows = [self._chunk_to_row(c) for c in chunks] + first_row = rows[0] + + table = self._qualified_table(collection_name) + # Assuming all chunks uniformly share the same metadata struct & embedding names config + cols = ", ".join(first_row.keys()) + placeholders = ", ".join(["%s"] * len(first_row)) + updates = ", ".join([f"{k}=EXCLUDED.{k}" for k in first_row.keys() if k != "id"]) + sql = f"INSERT INTO {table} ({cols}) VALUES ({placeholders}) ON CONFLICT (id) DO UPDATE SET {updates};" + + params_seq = [list(row.values()) for row in rows] + + with self._conn_context() as conn: + with conn.cursor() as cur: + # Use psycopg's executemany for batch inserts + cur.executemany(sql, params_seq) + + async def a_add( + self, chunk: Chunk | list[Chunk], collection_name: str | None = None + ): + if not collection_name: + raise ValueError("collection_name must be provided") + + chunks = [chunk] if isinstance(chunk, Chunk) else chunk + if not chunks: + return + + rows = [self._chunk_to_row(c) for c in chunks] + first_row = rows[0] + + table = self._qualified_table(collection_name) + cols = ", ".join(first_row.keys()) + placeholders = ", ".join(["%s"] * len(first_row)) + updates = ", ".join([f"{k}=EXCLUDED.{k}" for k in first_row.keys() if k != "id"]) + sql = f"INSERT INTO {table} ({cols}) VALUES ({placeholders}) ON CONFLICT (id) DO UPDATE SET {updates};" + + params_seq = [list(row.values()) for row in rows] + + async with await self._a_conn_context() as conn: + async with conn.cursor() as cur: + await cur.executemany(sql, params_seq) + + def _meta_table_name(self) -> str: + return f'"{self.schema}"."_datapizza_vectorstore_meta"' + + def _ensure_meta_table(self) -> None: + """Ensure the metadata table exists for storing collection configs.""" + with self._conn_context() as conn: + with conn.cursor() as cur: + self._execute( + cur, + f"CREATE TABLE IF NOT EXISTS {self._meta_table_name()} (" + "collection_name TEXT PRIMARY KEY, " + "config JSONB NOT NULL" + ");", + ) + + def _write_meta_config(self, collection_name: str, config: list[Dict[str, Any]]) -> None: + self._ensure_meta_table() + with self._conn_context() as conn: + with conn.cursor() as cur: + self._execute( + cur, + f"INSERT INTO {self._meta_table_name()} (collection_name, config) VALUES (%s, %s) " + "ON CONFLICT (collection_name) DO UPDATE SET config = EXCLUDED.config;", + (collection_name, Jsonb(config)), + ) + + def _read_meta_config(self, collection_name: str) -> Optional[list[Dict[str, Any]]]: + try: + with self._conn_context() as conn: + with conn.cursor(row_factory=dict_row) as cur: + self._execute( + cur, + f"SELECT config FROM {self._meta_table_name()} WHERE collection_name = %s;", + (collection_name,), + ) + row = cur.fetchone() + if not row: + return None + return row["config"] + except Exception: + return None + + def _load_collection_config_from_db( + self, collection_name: str + ) -> Optional[list[VectorConfig]]: + """Rebuild _collections cache by introspecting the Postgres schema. + + Uses a lightweight metadata table when available, falling back to + information_schema introspection. + """ + # Prefer stored config when available (fast + reliable) + meta = self._read_meta_config(collection_name) + if meta: + try: + configs = [VectorConfig(**c) for c in meta] + self._collections[collection_name] = configs + return configs + except Exception: + pass + + # a.atttypmod stores the vector dimension for pgvector (as 4 + 4*dim) + sql = ( + "SELECT a.attname, format_type(a.atttypid, a.atttypmod) as typ, a.atttypmod " + "FROM pg_attribute a " + "JOIN pg_class c ON a.attrelid = c.oid " + "JOIN pg_namespace n ON c.relnamespace = n.oid " + "WHERE n.nspname = %s AND c.relname = %s " + " AND a.attnum > 0 AND NOT a.attisdropped " + " AND format_type(a.atttypid, a.atttypmod) LIKE '%vector%';" + ) + schema = self.schema + table_name = collection_name + + try: + with self._conn_context() as conn: + with conn.cursor(row_factory=dict_row) as cur: + self._execute(cur, sql, (schema, table_name)) + rows = cur.fetchall() + except Exception: + # If we can't query (or the table doesn't exist), treat as missing. + return None + + vector_configs: list[VectorConfig] = [] + for row in rows: + typ = row.get("typ") + atttypmod = row.get("atttypmod") + + dim: Optional[int] = None + + # typ can be 'vector' or 'vector(N)' + if isinstance(typ, str) and typ.startswith("vector(") and typ.endswith(")"): + try: + dim = int(typ[len("vector(") : -1]) + except ValueError: + dim = None + elif isinstance(typ, str) and typ == "vector" and isinstance(atttypmod, int): + # pgvector uses atttypmod = 4 + 4*dim + try: + dim = (atttypmod - 4) // 4 + except Exception: + dim = None + + if dim is None: + continue + + vector_configs.append( + VectorConfig(name=row["attname"], dimensions=dim) + ) + + if vector_configs: + self._collections[collection_name] = vector_configs + return vector_configs + return None + + def list_collections(self) -> list[str]: + """Return the list of collections (tables) tracked by this vectorstore.""" + try: + with self._conn_context() as conn: + with conn.cursor(row_factory=dict_row) as cur: + self._execute( + cur, + f"SELECT collection_name FROM {self._meta_table_name()} ORDER BY collection_name;", + ) + rows = cur.fetchall() + return [row["collection_name"] for row in rows] + except Exception: + # If the meta table does not exist, fall back to introspection. + return [] + + def describe_collection(self, collection_name: str) -> dict[str, Any]: + """Return metadata about a collection (vector columns and dimensions).""" + configs = self._load_collection_config_from_db(collection_name) + if not configs: + raise ValueError(f"Collection '{collection_name}' does not exist") + + return { + "name": collection_name, + "schema": self.schema, + "vectors": [ + {"name": cfg.name or "embedding", "dimensions": cfg.dimensions} + for cfg in configs + ], + } + + def _get_vector_config( + self, collection_name: str, vector_name: Optional[str] + ) -> VectorConfig: + configs = self._collections.get(collection_name) + if not configs: + configs = self._load_collection_config_from_db(collection_name) + if not configs: + raise ValueError( + f"Collection '{collection_name}' is not registered or does not exist. " + "Call create_collection first." + ) + + dense_configs = [c for c in configs if c.format == EmbeddingFormat.DENSE] + if not dense_configs: + raise ValueError( + f"Collection '{collection_name}' has no dense vectors configured" + ) + + if vector_name: + for cfg in dense_configs: + if (cfg.name or "embedding") == vector_name: + return cfg + raise ValueError( + f"Vector name '{vector_name}' not found in collection '{collection_name}'" + ) + + if len(dense_configs) == 1: + return dense_configs[0] + + raise ValueError( + f"Collection '{collection_name}' has multiple vector columns; specify `vector_name`." + ) + + def _get_vector_column(self, collection_name: str, vector_name: Optional[str]) -> str: + cfg = self._get_vector_config(collection_name, vector_name) + return cfg.name or "embedding" + + def _row_to_chunk(self, row: dict, vector_name: Optional[str]) -> Chunk: + emb: List[Any] = [] + if vector_name and vector_name in row: + val = row[vector_name] + if val is not None: + emb.append(DenseEmbedding(name=vector_name, vector=list(val))) + + return Chunk( + id=str(row["id"]), + text=row.get("text", ""), + embeddings=emb, + metadata=row.get("metadata", {}) or {}, + ) + + def search( + self, + collection_name: str, + query_vector: list[float], + k: int = 10, + vector_name: str | None = None, + filters: Optional[Dict[str, Any]] = None, + **kwargs, + ) -> list[Chunk]: + if not collection_name: + raise ValueError("collection_name must be provided") + + cfg = self._get_vector_config(collection_name, vector_name) + col = cfg.name or "embedding" + table = self._qualified_table(collection_name) + + op = "<->" if cfg.distance == Distance.EUCLIDEAN else "<=>" + + filter_sql, filter_params = self._build_filter_clause(filters) + sql = ( + f"SELECT id, text, metadata, {col} FROM {table} " + f"{filter_sql} ORDER BY {col} {op} %s LIMIT %s;" + ) + + with self._conn_context() as conn: + with conn.cursor(row_factory=dict_row) as cur: + self._execute(cur, sql, [*filter_params, Vector(query_vector), k]) + rows = cur.fetchall() + + return [self._row_to_chunk(row, col) for row in rows] + + async def a_search( + self, + collection_name: str, + query_vector: list[float], + k: int = 10, + vector_name: str | None = None, + filters: Optional[Dict[str, Any]] = None, + **kwargs, + ) -> list[Chunk]: + if not collection_name: + raise ValueError("collection_name must be provided") + + col = self._get_vector_column(collection_name, vector_name) + table = self._qualified_table(collection_name) + + cfg = self._get_vector_config(collection_name, vector_name) + op = "<->" if cfg.distance == Distance.EUCLIDEAN else "<=>" + + filter_sql, filter_params = self._build_filter_clause(filters) + sql = ( + f"SELECT id, text, metadata, {col} FROM {table} " + f"{filter_sql} ORDER BY {col} {op} %s LIMIT %s;" + ) + + async with await self._a_conn_context() as conn: + async with conn.cursor(row_factory=dict_row) as cur: + await self._a_execute( + cur, sql, [*filter_params, Vector(query_vector), k] + ) + rows = await cur.fetchall() + + return [self._row_to_chunk(row, col) for row in rows] + + def retrieve( + self, + collection_name: str, + ids: list[str], + filters: Optional[Dict[str, Any]] = None, + **kwargs, + ) -> list[Chunk]: + if not collection_name: + raise ValueError("collection_name must be provided") + if not ids: + return [] + + table = self._qualified_table(collection_name) + filter_sql, filter_params = self._build_filter_clause(filters, prefix="AND") + sql = f"SELECT id, text, metadata FROM {table} WHERE id = ANY(%s) {filter_sql};" + + with self._conn_context() as conn: + with conn.cursor(row_factory=dict_row) as cur: + self._execute(cur, sql, [ids, *filter_params]) + rows = cur.fetchall() + + return [self._row_to_chunk(row, None) for row in rows] + + def remove( + self, + collection_name: str, + ids: list[str], + filters: Optional[Dict[str, Any]] = None, + **kwargs, + ): + if not collection_name: + raise ValueError("collection_name must be provided") + if not ids: + return + + table = self._qualified_table(collection_name) + filter_sql, filter_params = self._build_filter_clause(filters, prefix="AND") + sql = f"DELETE FROM {table} WHERE id = ANY(%s) {filter_sql};" + + with self._conn_context() as conn: + with conn.cursor() as cur: + self._execute(cur, sql, [ids, *filter_params]) + + def update( + self, + collection_name: str, + payload: dict, + points: list[str], + filters: Optional[Dict[str, Any]] = None, + **kwargs, + ): + """Update metadata for a set of points. + + Args: + collection_name: Target collection. + payload: Metadata changes to merge (JSONB merge semantics). + points: List of IDs to update. + """ + if not collection_name: + raise ValueError("collection_name must be provided") + if not points: + return + + table = self._qualified_table(collection_name) + filter_sql, filter_params = self._build_filter_clause(filters, prefix="AND") + sql = ( + f"UPDATE {table} SET metadata = metadata || %s WHERE id = ANY(%s) " + f"{filter_sql};" + ) + with self._conn_context() as conn: + with conn.cursor() as cur: + self._execute(cur, sql, [Jsonb(payload), points, *filter_params]) + + def _get_ops_class(self, distance: Distance) -> str: + return "vector_l2_ops" if distance == Distance.EUCLIDEAN else "vector_cosine_ops" + + def close(self): + """Close any open synchronous or asynchronous connections and pools.""" + if self._conn is not None: + try: + self._conn.close() + finally: + self._conn = None + if self._pool is not None: + try: + self._pool.close() + finally: + self._pool = None + + async def aclose(self): + """Close any open asynchronous connection (and sync connection) and pools.""" + if self._a_conn is not None: + try: + await self._a_conn.close() + finally: + self._a_conn = None + if self._a_pool is not None: + try: + await self._a_pool.close() + finally: + self._a_pool = None + + # Close sync connection too, for safety. + if self._conn is not None: + try: + self._conn.close() + finally: + self._conn = None + if self._pool is not None: + try: + self._pool.close() + finally: + self._pool = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + self.close() + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + await self.aclose() + + def drop_collection(self, collection_name: str, if_exists: bool = True): + """Drop the collection table from the database.""" + if not collection_name: + raise ValueError("collection_name must be provided") + + table = self._qualified_table(collection_name) + with self._conn_context() as conn: + with conn.cursor() as cur: + if if_exists: + self._execute(cur, f"DROP TABLE IF EXISTS {table};") + else: + self._execute(cur, f"DROP TABLE {table};") + + self._collections.pop(collection_name, None) + + async def a_drop_collection(self, collection_name: str, if_exists: bool = True): + """Async drop of the collection table.""" + if not collection_name: + raise ValueError("collection_name must be provided") + + table = self._qualified_table(collection_name) + async with await self._a_conn_context() as conn: + async with conn.cursor() as cur: + if if_exists: + await self._a_execute(cur, f"DROP TABLE IF EXISTS {table};") + else: + await self._a_execute(cur, f"DROP TABLE {table};") + + self._collections.pop(collection_name, None) + + def create_index( + self, + collection_name: str, + vector_name: Optional[str] = None, + method: str = "hnsw", + index_name: Optional[str] = None, + ivfflat_lists: Optional[int] = None, + ): + """Create an ANN index for the given vector column. + + Args: + collection_name: collection/table name. + vector_name: vector column name (required for multi-vector collections). + method: One of "hnsw" or "ivfflat". + index_name: Optional explicit index name; generated if missing. + ivfflat_lists: Required for ivfflat; ignored for hnsw. + """ + if not collection_name: + raise ValueError("collection_name must be provided") + + cfg = self._get_vector_config(collection_name, vector_name) + col = cfg.name or "embedding" + table = self._qualified_table(collection_name) + + ops = self._get_ops_class(cfg.distance) + method = method.lower() + if method not in {"hnsw", "ivfflat"}: + raise ValueError("method must be 'hnsw' or 'ivfflat'") + + if not index_name: + index_name = f"{collection_name}_{col}_{method}_idx" + + if method == "hnsw": + sql = f"CREATE INDEX IF NOT EXISTS {index_name} ON {table} USING hnsw ({col} {ops});" + else: + if ivfflat_lists is None: + raise ValueError("ivfflat requires ivfflat_lists to be provided") + sql = ( + f"CREATE INDEX IF NOT EXISTS {index_name} ON {table} " + f"USING ivfflat ({col} {ops}) WITH (lists = {ivfflat_lists});" + ) + + with self._conn_context() as conn: + with conn.cursor() as cur: + self._execute(cur, sql) + + async def a_create_index( + self, + collection_name: str, + vector_name: Optional[str] = None, + method: str = "hnsw", + index_name: Optional[str] = None, + ivfflat_lists: Optional[int] = None, + ): + """Async version of create_index.""" + if not collection_name: + raise ValueError("collection_name must be provided") + + cfg = self._get_vector_config(collection_name, vector_name) + col = cfg.name or "embedding" + table = self._qualified_table(collection_name) + + ops = self._get_ops_class(cfg.distance) + method = method.lower() + if method not in {"hnsw", "ivfflat"}: + raise ValueError("method must be 'hnsw' or 'ivfflat'") + + if not index_name: + index_name = f"{collection_name}_{col}_{method}_idx" + + if method == "hnsw": + sql = f"CREATE INDEX IF NOT EXISTS {index_name} ON {table} USING hnsw ({col} {ops});" + else: + if ivfflat_lists is None: + raise ValueError("ivfflat requires ivfflat_lists to be provided") + sql = ( + f"CREATE INDEX IF NOT EXISTS {index_name} ON {table} " + f"USING ivfflat ({col} {ops}) WITH (lists = {ivfflat_lists});" + ) + + async with await self._a_conn_context() as conn: + async with conn.cursor() as cur: + await self._a_execute(cur, sql) + + def list_indexes(self, collection_name: str) -> list[dict[str, Any]]: + """List known indexes for a collection.""" + if not collection_name: + raise ValueError("collection_name must be provided") + + sql = ( + "SELECT indexname, indexdef " + "FROM pg_indexes " + "WHERE schemaname = %s AND tablename = %s;" + ) + with self._conn_context() as conn: + with conn.cursor(row_factory=dict_row) as cur: + self._execute(cur, sql, (self.schema, collection_name)) + rows = cur.fetchall() + + return [ + {"name": r["indexname"], "definition": r["indexdef"]} for r in rows + ] + + def describe_index(self, index_name: str) -> dict[str, Any]: + """Return information about a specific index.""" + if not index_name: + raise ValueError("index_name must be provided") + + sql = ( + "SELECT schemaname, tablename, indexname, indexdef " + "FROM pg_indexes " + "WHERE schemaname = %s AND indexname = %s;" + ) + with self._conn_context() as conn: + with conn.cursor(row_factory=dict_row) as cur: + self._execute(cur, sql, (self.schema, index_name)) + row = cur.fetchone() + + if not row: + raise ValueError(f"Index '{index_name}' does not exist") + + definition = row["indexdef"] + + method = None + ivfflat_lists = None + if definition: + if "using hnsw" in definition.lower(): + method = "hnsw" + elif "using ivfflat" in definition.lower(): + method = "ivfflat" + # parse lists=... + import re + + m = re.search(r"lists\s*=\s*'?(\d+)'?", definition, re.IGNORECASE) + if m: + ivfflat_lists = int(m.group(1)) + + return { + "schema": row["schemaname"], + "table": row["tablename"], + "name": row["indexname"], + "definition": definition, + "method": method, + "ivfflat_lists": ivfflat_lists, + } + + def drop_index(self, index_name: str, if_exists: bool = True): + """Drop an index by name.""" + if not index_name: + raise ValueError("index_name must be provided") + + stmt = f"DROP INDEX {'IF EXISTS ' if if_exists else ''}{self.schema}.{index_name};" + with self._conn_context() as conn: + with conn.cursor() as cur: + self._execute(cur, stmt) + + async def a_drop_index(self, index_name: str, if_exists: bool = True): + """Async drop of an index.""" + if not index_name: + raise ValueError("index_name must be provided") + + stmt = f"DROP INDEX {'IF EXISTS ' if if_exists else ''}{self.schema}.{index_name};" + async with await self._a_conn_context() as conn: + async with conn.cursor() as cur: + await self._a_execute(cur, stmt) From e0bf39151d4c5b25fac3cd3028be2753dbfb6869 Mon Sep 17 00:00:00 2001 From: Daniele Date: Sun, 15 Mar 2026 10:43:19 +0100 Subject: [PATCH 6/8] Add PgVector vectorstore tests Add a comprehensive test suite for the PgVector vectorstore integration and unit behavior. --- .../tests/test_pgvector_additional.py | 283 ++++++++++++++++++ .../tests/test_pgvector_distance_operator.py | 59 ++++ .../tests/test_pgvector_integration.py | 244 +++++++++++++++ .../tests/test_pgvector_multi_vector.py | 49 +++ .../tests/test_pgvector_update_metadata.py | 40 +++ .../tests/test_pgvector_vectorstore.py | 15 + 6 files changed, 690 insertions(+) create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_additional.py create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_distance_operator.py create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_multi_vector.py create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_update_metadata.py create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_vectorstore.py diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_additional.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_additional.py new file mode 100644 index 00000000..3dbc9707 --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_additional.py @@ -0,0 +1,283 @@ +import asyncio +import uuid + +import pytest +from datapizza.core.vectorstore import VectorConfig +from datapizza.type import Chunk, DenseEmbedding + +from datapizza.vectorstores.pgvector import PgVectorVectorstore + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_pgvector_async_add_search(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "async_collection", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + chunk_id = str(uuid.uuid4()) + await store.a_add( + Chunk( + id=chunk_id, + text="async", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + ), + collection_name="async_collection", + ) + + results = await store.a_search( + collection_name="async_collection", + query_vector=[0.0, 0.0, 0.0, 0.0], + k=1, + ) + + assert len(results) == 1 + assert results[0].id == chunk_id + + +@pytest.mark.integration +def test_pgvector_update_multiple_records(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "update_batch", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + ids = [str(uuid.uuid4()) for _ in range(2)] + for i, chunk_id in enumerate(ids, start=1): + store.add( + Chunk( + id=chunk_id, + text=f"row{i}", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + metadata={"tag": f"initial{i}"}, + ), + collection_name="update_batch", + ) + + store.update( + collection_name="update_batch", + payload={"updated": True}, + points=ids, + ) + + results = store.retrieve("update_batch", ids=ids) + assert len(results) == 2 + for r in results: + assert r.metadata.get("updated") is True + assert r.metadata.get("tag") is not None + + +@pytest.mark.unit +def test_pgvector_search_in_empty_cache_raises(): + store = PgVectorVectorstore(dsn="postgresql://localhost:5432/test") + # Emulate case where create_collection was never called / cache not populated + store._collections.clear() + + with pytest.raises(ValueError): + store.search( + collection_name="missing", query_vector=[0.0, 0.0, 0.0, 0.0], k=1 + ) + + +@pytest.mark.integration +def test_pgvector_cache_rebuild_from_schema(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "cache_rebuild", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + # Clear the cache and ensure search still works by rebuilding from schema + store._collections.clear() + results = store.search( + collection_name="cache_rebuild", + query_vector=[0.0, 0.0, 0.0, 0.0], + k=1, + ) + + assert isinstance(results, list) + assert "cache_rebuild" in store._collections + + +@pytest.mark.integration +def test_pgvector_drop_collection(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "to_drop", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + store.drop_collection("to_drop") + + with pytest.raises(Exception): + store.search( + collection_name="to_drop", query_vector=[0.0, 0.0, 0.0, 0.0], k=1 + ) + + +@pytest.mark.integration +def test_pgvector_create_index(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "with_index", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + store.create_index("with_index", method="hnsw") + + conn = store._get_conn() + with conn.cursor() as cur: + cur.execute( + "SELECT indexname FROM pg_indexes WHERE schemaname=%s AND tablename=%s;", + (store.schema, "with_index"), + ) + indexes = [r[0] for r in cur.fetchall()] + + assert any("with_index" in idx for idx in indexes) + + +@pytest.mark.integration +def test_pgvector_context_manager_closes(pgvector_postgres): + with PgVectorVectorstore(dsn=pgvector_postgres) as store: + store.create_collection( + "cm", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + assert store._conn is None + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_pgvector_async_context_manager_closes(pgvector_postgres): + async with PgVectorVectorstore(dsn=pgvector_postgres) as store: + store.create_collection( + "cm_async", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + assert store._a_conn is None + + +@pytest.mark.integration +def test_pgvector_list_and_describe_collections(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "list_test", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + collections = store.list_collections() + assert "list_test" in collections + + desc = store.describe_collection("list_test") + assert desc["name"] == "list_test" + assert desc["schema"] == store.schema + assert "vectors" in desc and len(desc["vectors"]) == 1 + assert desc["vectors"][0]["name"] == "embedding" + assert isinstance(desc["vectors"][0]["dimensions"], (int, type(None))) + + +@pytest.mark.integration +def test_pgvector_index_listing_and_drop(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "index_test", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + # Create a named index + store.create_index( + "index_test", + method="hnsw", + index_name="idx_index_test_embedding_hnsw", + ) + + indexes = store.list_indexes("index_test") + assert any(i["name"] == "idx_index_test_embedding_hnsw" for i in indexes) + + info = store.describe_index("idx_index_test_embedding_hnsw") + assert info["table"] == "index_test" + + store.drop_index("idx_index_test_embedding_hnsw") + indexes_after = store.list_indexes("index_test") + assert all(i["name"] != "idx_index_test_embedding_hnsw" for i in indexes_after) + + +@pytest.mark.integration +@pytest.mark.integration +@pytest.mark.asyncio +async def test_pgvector_async_index_ivfflat(pgvector_postgres): + async with PgVectorVectorstore(dsn=pgvector_postgres) as store: + store.create_collection( + "index_test_ivfflat_async", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + await store.a_create_index( + "index_test_ivfflat_async", + method="ivfflat", + ivfflat_lists=8, + index_name="idx_index_test_ivfflat_async", + ) + + indexes = store.list_indexes("index_test_ivfflat_async") + assert any(i["name"] == "idx_index_test_ivfflat_async" for i in indexes) + + info = store.describe_index("idx_index_test_ivfflat_async") + assert info["table"] == "index_test_ivfflat_async" + assert info.get("method") == "ivfflat" + assert info.get("ivfflat_lists") == 8 + + await store.a_drop_index("idx_index_test_ivfflat_async") + indexes_after = store.list_indexes("index_test_ivfflat_async") + assert all(i["name"] != "idx_index_test_ivfflat_async" for i in indexes_after) + + +@pytest.mark.integration +def test_pgvector_full_integration_flow(pgvector_postgres): + """Full integration test: CRUD, search, index, list/describe, drop.""" + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "integration_smoke", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + create_index=True, + ) + + store.add( + Chunk( + id="smoke-1", + text="hello", + embeddings=[DenseEmbedding(name="embedding", vector=[0.1, 0.1, 0.1, 0.1])], + ), + collection_name="integration_smoke", + ) + + results = store.search( + collection_name="integration_smoke", + query_vector=[0.1, 0.1, 0.1, 0.1], + k=1, + ) + assert len(results) == 1 + assert results[0].id == "smoke-1" + + assert "integration_smoke" in store.list_collections() + + # Choose a non-pkey index to drop + indexes = store.list_indexes("integration_smoke") + idx_name = next(i["name"] for i in indexes if not i["name"].endswith("_pkey")) + + info = store.describe_index(idx_name) + assert info["table"] == "integration_smoke" + + store.drop_index(idx_name) + store.drop_collection("integration_smoke") diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_distance_operator.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_distance_operator.py new file mode 100644 index 00000000..48705adc --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_distance_operator.py @@ -0,0 +1,59 @@ +import pytest + +from datapizza.core.vectorstore import Distance, VectorConfig +from datapizza.type import Chunk, DenseEmbedding +from datapizza.vectorstores.pgvector import PgVectorVectorstore + + +class _FakeCursor: + def __init__(self): + self.queries = [] + self.params = [] + + def execute(self, sql, params=None): + self.queries.append(sql) + self.params.append(params) + + def fetchall(self): + return [] + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + return False + + +class _FakeConn: + def __init__(self): + self.cursor_obj = _FakeCursor() + + def cursor(self, *args, **kwargs): + return self.cursor_obj + + +@pytest.mark.parametrize( + "distance,expected_op", + [ + (Distance.COSINE, "<=>"), + (Distance.EUCLIDEAN, "<->"), + ], +) +def test_pgvector_search_operator(distance, expected_op): + """Verify that Distance maps to the correct pgvector operator.""" + + store = PgVectorVectorstore(dsn="postgresql://user:pass@localhost:5432/db") + store._collections["test"] = [VectorConfig(name="embedding", dimensions=4, distance=distance)] + + fake_conn = _FakeConn() + store._get_conn = lambda: fake_conn + + store.search( + collection_name="test", + query_vector=[0.0, 0.0, 0.0, 0.0], + k=1, + vector_name="embedding", + ) + + assert len(fake_conn.cursor_obj.queries) == 1 + assert expected_op in fake_conn.cursor_obj.queries[0] diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py new file mode 100644 index 00000000..62e7c0af --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py @@ -0,0 +1,244 @@ +import uuid + +import pytest +from datapizza.core.vectorstore import VectorConfig +from datapizza.type import Chunk, DenseEmbedding + +from datapizza.vectorstores.pgvector import PgVectorVectorstore + + +@pytest.mark.integration +def test_pgvector_collection_add_search(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "test_collection", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + chunk_id = str(uuid.uuid4()) + store.add( + Chunk( + id=chunk_id, + text="hello", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + metadata={"tag": "keep"}, + ), + collection_name="test_collection", + ) + + # Add a second chunk with different metadata to ensure filters work. + store.add( + Chunk( + id=str(uuid.uuid4()), + text="other", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + metadata={"tag": "skip"}, + ), + collection_name="test_collection", + ) + + results = store.search( + collection_name="test_collection", + query_vector=[0.0, 0.0, 0.0, 0.0], + k=10, + filters={"tag": "keep"}, + ) + + assert len(results) == 1 + assert results[0].id == chunk_id + assert results[0].text == "hello" + + +@pytest.mark.integration +def test_pgvector_collection_pooling_option(pgvector_postgres): + pytest.importorskip("psycopg_pool") + + store = PgVectorVectorstore(dsn=pgvector_postgres, use_pool=True) + + store.create_collection( + "test_collection_pool", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + chunk_id = str(uuid.uuid4()) + store.add( + Chunk( + id=chunk_id, + text="pool-test", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + ), + collection_name="test_collection_pool", + ) + + results = store.search( + collection_name="test_collection_pool", + query_vector=[0.0, 0.0, 0.0, 0.0], + k=1, + ) + + assert len(results) == 1 + assert results[0].id == chunk_id + + +def test_pgvector_collection_remove_and_retrieve(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "test_collection_2", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + chunk_id = str(uuid.uuid4()) + store.add( + Chunk( + id=chunk_id, + text="remove-me", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + ), + collection_name="test_collection_2", + ) + + # Verify retrieval works + retrieved = store.retrieve("test_collection_2", ids=[chunk_id]) + assert len(retrieved) == 1 + assert retrieved[0].id == chunk_id + + # Remove and verify gone + store.remove("test_collection_2", ids=[chunk_id]) + retrieved_after = store.retrieve("test_collection_2", ids=[chunk_id]) + assert retrieved_after == [] + + +def test_pgvector_update_and_remove_with_filters(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "test_collection_filters", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + # Add three chunks with varying metadata + chunk_a = Chunk( + id=str(uuid.uuid4()), + text="a", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + metadata={"tag": "keep", "group": "x"}, + ) + chunk_b = Chunk( + id=str(uuid.uuid4()), + text="b", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + metadata={"tag": "keep", "group": "y"}, + ) + chunk_c = Chunk( + id=str(uuid.uuid4()), + text="c", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + metadata={"tag": "skip", "group": "x"}, + ) + + store.add(chunk_a, collection_name="test_collection_filters") + store.add(chunk_b, collection_name="test_collection_filters") + store.add(chunk_c, collection_name="test_collection_filters") + + # Update only the "keep" group items to add a "status" field. + store.update( + "test_collection_filters", + payload={"status": "updated"}, + points=[chunk_a.id, chunk_b.id, chunk_c.id], + filters={"tag": "keep"}, + ) + + # Only the two matching items should have the new field. + results = store.search( + collection_name="test_collection_filters", + query_vector=[0.0, 0.0, 0.0, 0.0], + k=10, + filters={"status": "updated"}, + ) + assert {r.id for r in results} == {chunk_a.id, chunk_b.id} + + # Remove items with tag==skip + store.remove( + "test_collection_filters", + ids=[chunk_a.id, chunk_b.id, chunk_c.id], + filters={"tag": "skip"}, + ) + + remaining = store.search( + collection_name="test_collection_filters", + query_vector=[0.0, 0.0, 0.0, 0.0], + k=10, + ) + assert {r.id for r in remaining} == {chunk_a.id, chunk_b.id} + + +@pytest.mark.integration +def test_pgvector_batched_add_and_close(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres, use_pool=True) + + store.create_collection( + "test_batched", + vector_config=[VectorConfig(name="embedding", dimensions=2)], + ) + + # test batched insertion with executemany + chunks = [ + Chunk( + id=str(uuid.uuid4()), + text=f"batched {i}", + embeddings=[DenseEmbedding(name="embedding", vector=[float(i), float(i)])], + ) + for i in range(10) + ] + + store.add(chunks, collection_name="test_batched") + + results = store.search( + collection_name="test_batched", + query_vector=[0.0, 0.0], + k=20, + ) + assert len(results) == 10 + + # Ensure explicit close works properly without raising errors + assert store._pool is not None + store.close() + assert store._pool is None + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_pgvector_async_batched_add_and_aclose(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres, use_pool=True) + + # Sync create collection since a_create_collection is typically not required or missing in simple cases + # Assuming create_collection is sync only in this test context + store.create_collection( + "test_async_batched", + vector_config=[VectorConfig(name="embedding", dimensions=2)], + ) + + chunks = [ + Chunk( + id=str(uuid.uuid4()), + text=f"async batched {i}", + embeddings=[DenseEmbedding(name="embedding", vector=[float(i), float(i)])], + ) + for i in range(10) + ] + + await store.a_add(chunks, collection_name="test_async_batched") + + # Assuming a_search or search works + results = store.search( + collection_name="test_async_batched", + query_vector=[0.0, 0.0], + k=20, + ) + assert len(results) == 10 + + assert store._a_pool is not None + await store.aclose() + assert store._a_pool is None diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_multi_vector.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_multi_vector.py new file mode 100644 index 00000000..5085829f --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_multi_vector.py @@ -0,0 +1,49 @@ +import uuid + +import pytest +from datapizza.core.vectorstore import VectorConfig +from datapizza.type import Chunk, DenseEmbedding + +from datapizza.vectorstores.pgvector import PgVectorVectorstore + + +@pytest.mark.integration +def test_pgvector_multi_vector_requires_vector_name(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + # Create a collection with two vector fields. + store.create_collection( + "multi_vector", + vector_config=[ + VectorConfig(name="openai", dimensions=4), + VectorConfig(name="cohere", dimensions=4), + ], + ) + + chunk_id = str(uuid.uuid4()) + store.add( + Chunk( + id=chunk_id, + text="multi", + embeddings=[ + DenseEmbedding(name="openai", vector=[0.0, 0.0, 0.0, 0.0]), + DenseEmbedding(name="cohere", vector=[0.1, 0.1, 0.1, 0.1]), + ], + ), + collection_name="multi_vector", + ) + + # Query must specify vector name in multi-vector collections. + with pytest.raises(ValueError): + store.search(collection_name="multi_vector", query_vector=[0.0, 0.0, 0.0, 0.0]) + + # Works when specifying vector name + results = store.search( + collection_name="multi_vector", + query_vector=[0.0, 0.0, 0.0, 0.0], + vector_name="openai", + k=1, + ) + + assert len(results) == 1 + assert results[0].text == "multi" diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_update_metadata.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_update_metadata.py new file mode 100644 index 00000000..e2b3c2e4 --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_update_metadata.py @@ -0,0 +1,40 @@ +import uuid + +import pytest +from datapizza.core.vectorstore import VectorConfig +from datapizza.type import Chunk, DenseEmbedding + +from datapizza.vectorstores.pgvector import PgVectorVectorstore + + +@pytest.mark.integration +def test_pgvector_update_metadata_merges(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "update_metadata", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + chunk_id = str(uuid.uuid4()) + store.add( + Chunk( + id=chunk_id, + text="hello", + embeddings=[DenseEmbedding(name="embedding", vector=[0.0, 0.0, 0.0, 0.0])], + metadata={"foo": "bar"}, + ), + collection_name="update_metadata", + ) + + # Merge new metadata keys without dropping existing ones + store.update( + collection_name="update_metadata", + payload={"baz": "qux"}, + points=[chunk_id], + ) + + results = store.retrieve("update_metadata", ids=[chunk_id]) + assert len(results) == 1 + assert results[0].metadata["foo"] == "bar" + assert results[0].metadata["baz"] == "qux" diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_vectorstore.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_vectorstore.py new file mode 100644 index 00000000..0f356822 --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_vectorstore.py @@ -0,0 +1,15 @@ +import pytest + +from datapizza.vectorstores.pgvector import PgVectorVectorstore + + +def test_pgvector_vectorstore_imports(): + # Basic smoke test to validate the package is importable and the class is available. + store = PgVectorVectorstore(dsn="postgresql://localhost:5432/postgres") + assert store is not None + + +@pytest.mark.parametrize("schema", ["public", "test_schema"]) +def test_pgvector_vectorstore_schema(schema): + store = PgVectorVectorstore(dsn="postgresql://localhost:5432/postgres", schema=schema) + assert store.schema == schema From babd1c4c730c73407b43cc836173045b113c63b0 Mon Sep 17 00:00:00 2001 From: Daniele Date: Sun, 15 Mar 2026 10:44:34 +0100 Subject: [PATCH 7/8] Add pgvector vectorstore package and tests Introduce namespace support for datapizza.vectorstores and add a pgvector vectorstore package. --- .../datapizza/vectorstores/__init__.py | 11 ++++++++ .../vectorstores/pgvector/__init__.py | 5 ++++ .../vectorstores/pgvector/tests/conftest.py | 25 +++++++++++++++++++ docs/API Reference/Vectorstore/.pages | 4 +++ 4 files changed, 45 insertions(+) create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/__init__.py create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/__init__.py create mode 100644 datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/conftest.py create mode 100644 docs/API Reference/Vectorstore/.pages diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/__init__.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/__init__.py new file mode 100644 index 00000000..04da676f --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/__init__.py @@ -0,0 +1,11 @@ +# Namespace package support for datapizza.vectorstores + +from __future__ import annotations + +# This file exists so that `datapizza.vectorstores` can be imported both from +# the core datapizza package and from separately-installed vectorstore plugins. +# +# It uses pkgutil.extend_path to support namespace package behavior across +# multiple distributions. + +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/__init__.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/__init__.py new file mode 100644 index 00000000..9efdeddf --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/__init__.py @@ -0,0 +1,5 @@ +"""pgvector vectorstore implementation for datapizza-ai.""" + +from .pgvector_vectorstore import PgVectorVectorstore + +__all__ = ["PgVectorVectorstore"] diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/conftest.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/conftest.py new file mode 100644 index 00000000..a1ac5401 --- /dev/null +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/conftest.py @@ -0,0 +1,25 @@ +import pytest + + +@pytest.fixture(scope="session") +def pgvector_postgres(): + """Start a Postgres container with pgvector enabled.""" + try: + import psycopg + from testcontainers.postgres import PostgresContainer + + with PostgresContainer("pgvector/pgvector:pg16") as pg: + dsn = pg.get_connection_url() + # Normalize SQLAlchemy style DSN to libpq format (psycopg expects it). + if dsn.startswith("postgresql+psycopg2://"): + dsn = dsn.replace("postgresql+psycopg2://", "postgresql://", 1) + + # Ensure pgvector extension is installed in the database. + with psycopg.connect(dsn) as conn: + with conn.cursor() as cur: + cur.execute("CREATE EXTENSION IF NOT EXISTS vector;") + yield dsn + except Exception as e: + pytest.skip( + f"Skipping integration tests because docker/testcontainers is unavailable: {e}" + ) diff --git a/docs/API Reference/Vectorstore/.pages b/docs/API Reference/Vectorstore/.pages new file mode 100644 index 00000000..9d039d45 --- /dev/null +++ b/docs/API Reference/Vectorstore/.pages @@ -0,0 +1,4 @@ +nav: + - milvus_vectorstore.md + - qdrant_vectorstore.md + - pgvector_vectorstore.md From e4f82383496288d043531428ce1c756702e8d73c Mon Sep 17 00:00:00 2001 From: Daniele Date: Tue, 17 Mar 2026 11:12:45 +0100 Subject: [PATCH 8/8] Validate and cache pgvector collection schema Normalize and verify vector collection schemas when creating collections in the PgVector vectorstore. --- .../datapizza/core/vectorstore/vectorstore.py | 2 +- .../pgvector/pgvector_vectorstore.py | 49 ++++++++++++++++--- .../tests/test_pgvector_integration.py | 16 ++++++ 3 files changed, 59 insertions(+), 8 deletions(-) diff --git a/datapizza-ai-core/datapizza/core/vectorstore/vectorstore.py b/datapizza-ai-core/datapizza/core/vectorstore/vectorstore.py index b8757937..26bdab63 100644 --- a/datapizza-ai-core/datapizza/core/vectorstore/vectorstore.py +++ b/datapizza-ai-core/datapizza/core/vectorstore/vectorstore.py @@ -43,7 +43,7 @@ async def a_add( pass @abstractmethod - def update(self, collection_name: str, payload: dict, points: list[int], **kwargs): + def update(self, collection_name: str, payload: dict, points: list[str | int], **kwargs): pass @abstractmethod diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py index 4862341d..03820aaf 100644 --- a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/pgvector_vectorstore.py @@ -303,11 +303,8 @@ def create_collection( table = self._qualified_table(collection_name) - # Persist config locally for name inference - self._collections[collection_name] = vector_config - - # Persist config to the metadata table so we can rebuild the cache later - meta = [ + # Build a normalized schema description for the provided config. + requested_meta = [ { "name": cfg.name, "dimensions": cfg.dimensions, @@ -316,7 +313,21 @@ def create_collection( } for cfg in vector_config ] - self._write_meta_config(collection_name, meta) + + # If the collection was already created, ensure the requested schema matches. + existing_meta = self._read_meta_config(collection_name) + if existing_meta is not None: + if existing_meta != requested_meta: + raise ValueError( + f"Collection '{collection_name}' already exists with a different vector schema. " + "Call drop_collection() first if you want to recreate it." + ) + # Cache the existing config for name inference. + self._collections[collection_name] = [VectorConfig(**m) for m in existing_meta] + return + + # Cache the requested config for name inference. + self._collections[collection_name] = vector_config # Build schema cols = [ @@ -341,6 +352,30 @@ def create_collection( self._execute(cur, f"CREATE SCHEMA IF NOT EXISTS \"{self.schema}\";") self._execute(cur, f"CREATE TABLE IF NOT EXISTS {table} ({', '.join(cols)});") + # Validate that the created table matches the requested config if possible. + # This is best-effort: some pgvector / Postgres versions may not expose the + # vector type metadata in a way we can reliably introspect. + actual_config = self._load_collection_config_from_db(collection_name) + if actual_config is None: + log.warning( + "Unable to introspect collection '%s' after creation; continuing without schema verification.", + collection_name, + ) + else: + def _schema_key(cfg: VectorConfig) -> tuple[str | None, int | None]: + return (cfg.name, cfg.dimensions) + + requested_keys = {_schema_key(c) for c in vector_config} + actual_keys = {_schema_key(c) for c in actual_config} + if requested_keys != actual_keys: + raise ValueError( + f"Collection '{collection_name}' exists with a different schema: " + f"requested={requested_keys} actual={actual_keys}" + ) + + # Persist config to the metadata table for faster subsequent loads. + self._write_meta_config(collection_name, requested_meta) + if create_index: self.create_index( collection_name=collection_name, @@ -727,7 +762,7 @@ def update( self, collection_name: str, payload: dict, - points: list[str], + points: list[str | int], filters: Optional[Dict[str, Any]] = None, **kwargs, ): diff --git a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py index 62e7c0af..52f5f33e 100644 --- a/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py +++ b/datapizza-ai-vectorstores/datapizza-ai-vectorstores-pgvector/datapizza/vectorstores/pgvector/tests/test_pgvector_integration.py @@ -50,6 +50,22 @@ def test_pgvector_collection_add_search(pgvector_postgres): assert results[0].text == "hello" +@pytest.mark.integration +def test_pgvector_collection_create_collection_schema_mismatch_raises(pgvector_postgres): + store = PgVectorVectorstore(dsn=pgvector_postgres) + + store.create_collection( + "test_collection_mismatch", + vector_config=[VectorConfig(name="embedding", dimensions=4)], + ) + + with pytest.raises(ValueError): + store.create_collection( + "test_collection_mismatch", + vector_config=[VectorConfig(name="embedding", dimensions=8)], + ) + + @pytest.mark.integration def test_pgvector_collection_pooling_option(pgvector_postgres): pytest.importorskip("psycopg_pool")