wenjin272 commented on code in PR #1143: URL: https://github.com/apache/flink-agents/pull/1143#discussion_r4409528303
########## integrations/vector-stores/pgvector/src/main/java/org/apache/flink/agents/integrations/vectorstores/pgvector/PgVectorVectorStore.java: ########## @@ -0,0 +1,1684 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.integrations.vectorstores.pgvector; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.vectorstores.BaseVectorStore; +import org.apache.flink.agents.api.vectorstores.CollectionManageableVectorStore; +import org.apache.flink.agents.api.vectorstores.Document; +import org.postgresql.ds.PGSimpleDataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.sql.Array; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.regex.Pattern; + +/** + * PostgreSQL <a href="https://github.com/pgvector/pgvector";>pgvector</a> backed implementation of a + * vector store. + * + * <p>Each collection is a table with a text primary key, the document content, the metadata as + * {@code jsonb}, and the embedding as a {@code vector(dims)} column. Similarity search orders rows + * by the pgvector distance operator of the configured metric; equality filters from the unified + * filter DSL compile to a single {@code jsonb} containment predicate ({@code metadata @> '{...}'}), + * which the default GIN index on the metadata column serves. + * + * <p>Configuration is provided through {@link ResourceDescriptor} arguments: + * + * <ul> + * <li>{@code uri} (optional): JDBC URL, default {@code + * jdbc:postgresql://localhost:5432/postgres}. Alternatively {@code host}, {@code port} and + * {@code database} build the URL. + * <li>{@code username}, {@code password} (optional): connection credentials. + * <li>{@code connect_timeout_s} (optional): seconds to wait for a TCP connection and login. + * {@code socket_timeout_s} (optional): seconds to wait for a reply to any statement before + * the connection is dropped. When omitted, the driver defaults (10 seconds to connect, no + * socket limit) or the {@code connectTimeout}/{@code loginTimeout}/{@code socketTimeout} + * parameters of {@code uri} apply. Set a socket timeout when the job must fail over rather + * than wait on a hung server; keep it above the longest index build you expect from {@code + * createCollectionIfNotExists}. + * <li>{@code collection} (optional): default table name, default {@link #DEFAULT_COLLECTION}; + * {@code collection_name} and {@code index} are accepted as aliases. {@code schema} + * (optional) selects the PostgreSQL schema, default {@code public}. + * <li>{@code id_field}, {@code content_field}, {@code metadata_field}, {@code vector_field} + * (optional): column names. + * <li>{@code dims} (optional): vector dimensionality, default {@link #DEFAULT_DIMENSION}. + * <li>{@code metric_type} (optional): {@code COSINE} (default), {@code L2} or {@code IP}. + * <li>{@code index_type} (optional): {@code HNSW} (default), {@code IVFFLAT} or {@code NONE}. + * <li>{@code index_params} (optional): index storage parameters, for example {@code m} and {@code + * ef_construction} for HNSW or {@code lists} for IVFFlat. + * <li>{@code create_extension} (optional): run {@code CREATE EXTENSION IF NOT EXISTS vector} + * before creating a table, default {@code true}. The statement is a no-op once the extension + * is installed; set {@code false} to skip it entirely, for example where DDL is audited or + * the role may not issue it. + * <li>{@code iterative_scan} (optional): {@code RELAXED_ORDER} (default), {@code STRICT_ORDER} or + * {@code OFF}. On pgvector 0.8 or later every similarity search runs with this iterative scan + * mode ({@code SET LOCAL} in the search's own transaction block), so that an HNSW or IVFFlat + * index keeps searching until {@code limit} rows satisfy the filter instead of returning only + * the matches among its first {@code ef_search} candidates. {@code OFF} leaves the server's + * settings alone. A per-query {@code iterative_scan} argument overrides it for that query. + * </ul> + * + * <p>Names are quoted as given, so they are matched case-sensitively; a table created without + * quotes has a lower-case name in PostgreSQL and must be referenced in lower case. HNSW and IVFFlat + * indexes support at most {@value #MAX_INDEXED_DIMENSIONS} dimensions, so {@code + * createCollectionIfNotExists} rejects wider vectors unless {@code index_type} is {@code NONE}; + * existing tables of any width can be used. Each store instance holds one connection and serializes + * its operations. + * + * <p>Score semantics: {@link Document#getScore()} holds the cosine similarity for {@code COSINE} + * and the inner product for {@code IP} (higher is better), and the Euclidean distance for {@code + * L2} (lower is better). Results are always returned nearest first, also under a relaxed iterative + * scan: the hits are re-sorted before they are returned. + * + * <p>Example usage: + * + * <pre>{@code + * ResourceDescriptor desc = ResourceDescriptor.Builder + * .newBuilder(PgVectorVectorStore.class.getName()) + * .addInitialArgument("embedding_model", "textEmbedder") + * .addInitialArgument("uri", "jdbc:postgresql://localhost:5432/postgres") + * .addInitialArgument("username", "postgres") + * .addInitialArgument("password", "postgres") + * .addInitialArgument("collection", "my_documents") + * .addInitialArgument("dims", 768) + * .addInitialArgument("metric_type", "COSINE") + * .build(); + * }</pre> + */ +public class PgVectorVectorStore extends BaseVectorStore + implements CollectionManageableVectorStore { + + /** Default table name used when {@code collection} and its aliases are omitted. */ + public static final String DEFAULT_COLLECTION = "flink_agents_pgvector_collection"; + /** Default PostgreSQL schema. */ + public static final String DEFAULT_SCHEMA = "public"; + /** Defaults for {@code host}, {@code port} and {@code database} when {@code uri} is omitted. */ + public static final String DEFAULT_HOST = "localhost"; + + public static final int DEFAULT_PORT = 5432; + public static final String DEFAULT_DATABASE = "postgres"; + /** Default JDBC URL, built from the three defaults above. */ + public static final String DEFAULT_URI = jdbcUrl(DEFAULT_HOST, DEFAULT_PORT, DEFAULT_DATABASE); + /** Default primary key column name. */ + public static final String DEFAULT_ID_FIELD = "id"; + /** Default column name used to store document content. */ + public static final String DEFAULT_CONTENT_FIELD = "content"; + /** Default {@code jsonb} column name used to store document metadata. */ + public static final String DEFAULT_METADATA_FIELD = "metadata"; + /** Default {@code vector} column name on which similarity search is executed. */ + public static final String DEFAULT_VECTOR_FIELD = "embedding"; + /** Default vector dimensionality used when {@code dims} is not provided. */ + public static final int DEFAULT_DIMENSION = 768; + /** First pgvector version with {@code hnsw.iterative_scan} / {@code ivfflat.iterative_scan}. */ + static final int[] ITERATIVE_SCAN_MIN_VERSION = {0, 8}; + + private static final Logger LOG = LoggerFactory.getLogger(PgVectorVectorStore.class); + + /** Distance metrics supported by pgvector for {@code vector} columns. */ + public enum MetricType { + /** Cosine distance operator {@code <=>}; the score is the cosine similarity. */ + COSINE("<=>", "vector_cosine_ops"), + /** Euclidean distance operator {@code <->}; the score is the distance. */ + L2("<->", "vector_l2_ops"), + /** Negative inner product operator {@code <#>}; the score is the inner product. */ + IP("<#>", "vector_ip_ops"); + + private final String operator; + private final String operatorClass; + + MetricType(String operator, String operatorClass) { + this.operator = operator; + this.operatorClass = operatorClass; + } + + /** Returns the pgvector distance operator of this metric. */ + public String getOperator() { + return operator; + } + + /** Returns the pgvector operator class used when indexing for this metric. */ + public String getOperatorClass() { + return operatorClass; + } + + /** Converts the raw operator result into the documented score of this metric. */ + float toScore(double distance) { + switch (this) { + case COSINE: + return (float) (1.0 - distance); + case IP: + return (float) -distance; + default: + return (float) distance; + } + } + } + + /** Vector index types supported for created tables. */ + public enum IndexType { + HNSW, + IVFFLAT, + /** Do not create a vector index; searches scan the table. */ + NONE + } + + /** Iterative index scan modes applied to similarity searches (pgvector 0.8+). */ + public enum IterativeScan { + /** Do not change the server's scan settings. */ + OFF, + /** Keep scanning until enough rows match; results may be slightly out of order. */ + RELAXED_ORDER, + /** Keep scanning until enough rows match, in exact distance order (HNSW only). */ + STRICT_ORDER; + + String sqlValue() { + return name().toLowerCase(Locale.ROOT); + } + } + + private static final Pattern IDENTIFIER = Pattern.compile("[A-Za-z_][A-Za-z0-9_]*"); + /** PostgreSQL silently truncates longer identifiers, which would make names collide. */ + static final int MAX_IDENTIFIER_LENGTH = 63; + + /** + * Longest indexed column name for which {@link #indexName(String, String)} can still derive a + * unique index name for any table: {@code 63 - 8 (hash) - 1 (separator) - 5 ("_" and "_idx") - + * 1 (table prefix)}. + */ + static final int MAX_INDEXED_COLUMN_LENGTH = + MAX_IDENTIFIER_LENGTH - 8 - 1 - "_".length() - "_idx".length() - 1; + + /** Advisory lock key serialising {@code CREATE EXTENSION} across all tables. */ + private static final String EXTENSION_LOCK_KEY = "flink_agents_pgvector.extension"; + /** pgvector's HNSW and IVFFlat indexes reject {@code vector} columns wider than this. */ + static final int MAX_INDEXED_DIMENSIONS = 2000; + + private static final Pattern INDEX_PARAM_VALUE = Pattern.compile("[A-Za-z0-9_]+"); + + private final ObjectMapper mapper = new ObjectMapper(); + + /** JDBC URL. */ + private final String uri; + /** JDBC URL with any {@code user}/{@code password} query parameters removed. */ + private final String publicUri; + /** Credentials; never exposed through {@link #getStoreKwargs()}. */ + private final @Nullable String username; + + private final @Nullable String password; + /** Explicit timeouts; {@code null} leaves the driver's (or the URL's) setting in place. */ + private final @Nullable Integer connectTimeoutSeconds; + + private final @Nullable Integer socketTimeoutSeconds; + /** PostgreSQL schema that holds the tables. */ + private final String schema; + /** Default table name used when a per-call collection is not supplied. */ + private final String defaultCollection; + + private final String idField; + private final String contentField; + private final String metadataField; + private final String vectorField; + private final int dims; + private final MetricType metricType; + private final IndexType indexType; + private final Map<String, Object> indexParams; + private final boolean createExtension; + private final IterativeScan iterativeScan; + + /** Whether the connected server's pgvector supports iterative scans; resolved lazily. */ + private @Nullable Boolean iterativeScanSupported; + + /** + * Data source used to open connections. It carries the URL, credentials and timeouts, and its + * class initialiser loads {@code org.postgresql.Driver}, which registers the driver with {@code + * DriverManager} before {@code getConnection} goes through it; that global registration is why + * the docs recommend Flink's {@code lib/} for the dist jar on long-lived session clusters. + */ + private @Nullable PGSimpleDataSource dataSource; + /** + * One cached connection per store instance, revalidated and reopened when it broke. Operations + * are {@code synchronized}, so concurrent actions on the same instance are serialized on it + * rather than paying a connection handshake per call; parallelism comes from Flink subtasks, + * each of which owns its own store instance. + */ + private @Nullable Connection connection; + + /** Set by {@link #close()}; a closed store refuses to reconnect instead of leaking one. */ + private boolean closed; + + /** + * Tables this instance has already created or verified, keyed by table and the expected + * width/metric/index, so a repeated "ensure" call with the same expectations is a single + * existence lookup. {@link #deleteCollection} forgets the table again. + */ + private final Set<String> verifiedTables = new HashSet<>(); + + /** + * Creates a new {@code PgVectorVectorStore} from the provided descriptor and resource resolver. + * The JDBC connection is opened lazily on first use. + */ + public PgVectorVectorStore(ResourceDescriptor descriptor, ResourceContext resourceContext) { + super(descriptor, resourceContext); + + this.uri = resolveUri(descriptor); + this.publicUri = stripCredentials(this.uri); + this.username = stringArg(descriptor, "username", null); + this.password = stringArg(descriptor, "password", null); + this.connectTimeoutSeconds = optionalSeconds(descriptor, "connect_timeout_s"); + this.socketTimeoutSeconds = optionalSeconds(descriptor, "socket_timeout_s"); + this.schema = identifierArg(descriptor, "schema", DEFAULT_SCHEMA); + this.defaultCollection = + identifier( + "collection", + stringArg( + descriptor, + "collection", + stringArg( + descriptor, + "collection_name", + stringArg(descriptor, "index", DEFAULT_COLLECTION)))); + this.idField = identifierArg(descriptor, "id_field", DEFAULT_ID_FIELD); + this.contentField = identifierArg(descriptor, "content_field", DEFAULT_CONTENT_FIELD); + this.metadataField = + indexedColumn( + "metadata_field", + identifierArg(descriptor, "metadata_field", DEFAULT_METADATA_FIELD)); + this.vectorField = + indexedColumn( + "vector_field", + identifierArg(descriptor, "vector_field", DEFAULT_VECTOR_FIELD)); + this.dims = positive("dims", intArg(descriptor, "dims", DEFAULT_DIMENSION)); + this.metricType = + enumArg( + MetricType.class, + "metric_type", + stringArg(descriptor, "metric_type", MetricType.COSINE.name())); + this.indexType = + enumArg( + IndexType.class, + "index_type", + stringArg(descriptor, "index_type", IndexType.HNSW.name())); + this.indexParams = indexParams(mapArg(descriptor, "index_params")); + this.createExtension = booleanArg(descriptor, "create_extension", true); + this.iterativeScan = + enumArg( + IterativeScan.class, + "iterative_scan", + stringArg( + descriptor, "iterative_scan", IterativeScan.RELAXED_ORDER.name())); + } + + /** + * Closes the cached connection. Failures are logged, never thrown, so teardown cannot mask the + * failure that caused it. + */ + @Override + public synchronized void close() { + this.closed = true; + discardConnection(); + this.dataSource = null; + } + + /** + * Returns default store-level arguments collected from the descriptor. Credentials are not + * included. + */ + @Override + public Map<String, Object> getStoreKwargs() { + Map<String, Object> kwargs = new HashMap<>(); + kwargs.put("uri", this.publicUri); + kwargs.put("schema", this.schema); + kwargs.put("collection", this.defaultCollection); + kwargs.put("index", this.defaultCollection); + kwargs.put("id_field", this.idField); + kwargs.put("content_field", this.contentField); + kwargs.put("metadata_field", this.metadataField); + kwargs.put("vector_field", this.vectorField); + kwargs.put("dims", this.dims); + kwargs.put("metric_type", this.metricType.name()); + kwargs.put("index_type", this.indexType.name()); + kwargs.put("index_params", new LinkedHashMap<>(this.indexParams)); + kwargs.put("create_extension", this.createExtension); + kwargs.put("iterative_scan", this.iterativeScan.name()); + return kwargs; + } + + /** + * Creates the table for the given collection if it does not exist, together with the vector + * index for the configured metric and a GIN index on the metadata column. {@code dims}, {@code + * metric_type}, {@code index_type} and {@code index_params} in {@code kwargs} override the + * descriptor defaults for a new table. + */ + @Override + public synchronized void createCollectionIfNotExists( + String name, @Nullable Map<String, Object> options) throws Exception { + String table = identifier("collection", name); + TableSpec spec = tableSpec(options); + int dimension = spec.dimension; + MetricType metric = spec.metric; + IndexType index = spec.index; + Map<String, Object> params = spec.indexParams; + + // "Ensure" always looks at the catalog, so a table dropped behind our back is recreated. + // The common case, an existing table, is that one autocommit lookup and no lock; the + // schema inspection below runs once per table and expectation for this instance. + String verificationKey = table + "|" + dimension + "|" + metric + "|" + index; + boolean existed; + try { + existed = tableExists(connection(), table); + } catch (SQLException e) { + discardIfBroken(e); + throw e; + } + if (existed && this.verifiedTables.contains(verificationKey)) { + return; + } + if (!existed) { + // Only a table about to be created must fit its index; existing tables of any width + // can be used (with index_type NONE, or with whatever index they already carry). + checkIndexedDimensions(dimension, index); + // One transaction under advisory locks: parallel subtasks creating the extension or + // the same table would otherwise race past IF NOT EXISTS and fail on the catalog's + // unique constraints. The extension lock is always taken first, so lock order is + // consistent. + inTransaction( + connection -> { + try (Statement statement = connection.createStatement(); + PreparedStatement lock = + connection.prepareStatement( + "SELECT pg_advisory_xact_lock(hashtext(?))")) { + if (this.createExtension) { + lock.setString(1, EXTENSION_LOCK_KEY); + lock.execute(); + statement.execute("CREATE EXTENSION IF NOT EXISTS vector"); + // The extension (and hence the answer) may have just appeared. + this.iterativeScanSupported = null; + } + lock.setString(1, this.schema + "." + table); + lock.execute(); + // Another subtask may have created it between the check and the + // lock; IF NOT EXISTS below handles that, and reporting "created" for + // it only means the schema check runs, which it does anyway. + createTableAndIndexes( + statement, table, dimension, metric, index, params); + } + return null; + }); + } + // IF NOT EXISTS only compares names: an existing table may have been created for another + // width or metric, and CREATE INDEX IF NOT EXISTS is skipped when an unrelated relation + // already carries the index name. Say so instead of silently serving unindexed queries. + try { + for (String mismatch : + schemaMismatches(connection(), table, dimension, metric, index)) { + LOG.warn( + "Collection {} ({}): {}", + qualified(table), + existed ? "already existed" : "just created", + mismatch); + } + this.verifiedTables.add(verificationKey); + } catch (SQLException e) { + discardIfBroken(e); + throw e; + } + } + + /** + * Whether a plain or partitioned table exists by that name. A view, sequence or other relation + * occupying the name is reported as an error rather than silently treated as absent, because + * {@code CREATE TABLE IF NOT EXISTS} would skip it and the index creation would then fail. + */ + private boolean tableExists(Connection connection, String table) throws SQLException { + try (PreparedStatement statement = + connection.prepareStatement( + "SELECT c.relkind FROM pg_class c" + + " JOIN pg_namespace n ON n.oid = c.relnamespace" + + " WHERE n.nspname = ? AND c.relname = ?")) { + statement.setString(1, this.schema); + statement.setString(2, table); + try (ResultSet rows = statement.executeQuery()) { + if (!rows.next()) { + return false; + } + String kind = rows.getString(1); + if ("r".equals(kind) || "p".equals(kind)) { + return true; + } + throw new IllegalStateException( + qualified(table) + + " exists but is not a table (relkind '" + + kind + + "'); choose another collection name."); + } + } + } + + private void createTableAndIndexes( + Statement statement, + String table, + int dimension, + MetricType metric, + IndexType index, + Map<String, Object> params) + throws SQLException { + statement.execute( + "CREATE TABLE IF NOT EXISTS " + + qualified(table) + + " (" + + quote(this.idField) + + " TEXT PRIMARY KEY, " + + quote(this.contentField) + + " TEXT NOT NULL, " + + quote(this.metadataField) + + " JSONB, " + + quote(this.vectorField) + + " VECTOR(" + + dimension + + "))"); + statement.execute( + "CREATE INDEX IF NOT EXISTS " + + quote(indexName(table, this.metadataField)) + + " ON " + + qualified(table) + + " USING GIN (" + + quote(this.metadataField) + + ")"); + if (index != IndexType.NONE) { + statement.execute( + "CREATE INDEX IF NOT EXISTS " + + quote(indexName(table, this.vectorField)) + + " ON " + + qualified(table) + + " USING " + + index.name().toLowerCase(Locale.ROOT) + + " (" + + quote(this.vectorField) + + " " + + metric.getOperatorClass() + + ")" + + withClause(params)); + } + } + + /** + * Compares an existing collection's vector column and index with this store's expectations and + * returns one message per mismatch (empty when everything matches), for callers that prefer to + * fail fast over reading the warnings {@link #createCollectionIfNotExists} logs. {@code + * options} are interpreted as in {@link #createCollectionIfNotExists}. + */ + public synchronized List<String> schemaMismatches( + String name, @Nullable Map<String, Object> options) throws IOException { + String table = identifier("collection", name); + TableSpec spec = tableSpec(options); + try { + return schemaMismatches(connection(), table, spec.dimension, spec.metric, spec.index); + } catch (SQLException e) { + discardIfBroken(e); + throw new IOException("Failed to inspect " + qualified(table) + ".", e); + } + } + + /** What a collection's table should look like: descriptor defaults overridden per call. */ + private static final class TableSpec { + final int dimension; + final MetricType metric; + final IndexType index; + final Map<String, Object> indexParams; + + TableSpec(int dimension, MetricType metric, IndexType index, Map<String, Object> params) { + this.dimension = dimension; + this.metric = metric; + this.index = index; + this.indexParams = params; + } + } + + /** + * Resolves {@code dims}, {@code metric_type}, {@code index_type} and {@code index_params} from + * per-call options over the descriptor defaults; {@code null} options (or values) fall back. + */ + private TableSpec tableSpec(@Nullable Map<String, Object> options) { + Map<String, Object> kwargs = options == null ? Collections.emptyMap() : options; + int dimension = positive("dims", intFromMap(kwargs, "dims", this.dims)); + MetricType metric = enumFromMap(MetricType.class, kwargs, "metric_type", this.metricType); + IndexType index = enumFromMap(IndexType.class, kwargs, "index_type", this.indexType); + Object params = kwargs.get("index_params"); + return new TableSpec( + dimension, + metric, + index, + params == null ? this.indexParams : indexParams(toMap("index_params", params))); + } + + private List<String> schemaMismatches( + Connection connection, String table, int dimension, MetricType metric, IndexType index) + throws SQLException { + List<String> mismatches = new ArrayList<>(); + try (PreparedStatement statement = + connection.prepareStatement( + "SELECT t.typname, a.atttypmod FROM pg_attribute a" + + " JOIN pg_class c ON c.oid = a.attrelid" + + " JOIN pg_namespace n ON n.oid = c.relnamespace" + + " JOIN pg_type t ON t.oid = a.atttypid" + + " WHERE n.nspname = ? AND c.relname = ? AND a.attname = ?")) { + statement.setString(1, this.schema); + statement.setString(2, table); + statement.setString(3, this.vectorField); + try (ResultSet rows = statement.executeQuery()) { + if (!rows.next()) { + mismatches.add( + "table or column " + quote(this.vectorField) + " does not exist"); + return mismatches; + } else if (!"vector".equals(rows.getString(1))) { + mismatches.add( + "column " + + quote(this.vectorField) + + " has type " + + rows.getString(1) + + ", expected vector"); + } else if (rows.getInt(2) != dimension) { + // pgvector stores the declared width as the type modifier (-1 = unbounded). + mismatches.add( + "column " + + quote(this.vectorField) + + " is vector(" + + (rows.getInt(2) < 0 ? "" : rows.getInt(2)) + + "), expected vector(" + + dimension + + ")"); + } + } + } + if (index == IndexType.NONE) { + return mismatches; + } + List<String> vectorIndexes = new ArrayList<>(); + boolean usable = false; + try (PreparedStatement statement = + connection.prepareStatement( + "SELECT am.amname, oc.opcname FROM pg_index i" + + " JOIN pg_class c ON c.oid = i.indrelid" + + " JOIN pg_namespace n ON n.oid = c.relnamespace" + + " JOIN pg_class ic ON ic.oid = i.indexrelid" + + " JOIN pg_am am ON am.oid = ic.relam" + + " JOIN pg_opclass oc ON oc.oid = i.indclass[0]" + + " JOIN pg_attribute a ON a.attrelid = c.oid" + + " AND a.attnum = i.indkey[0]" + + " WHERE n.nspname = ? AND c.relname = ? AND a.attname = ?" + + " AND am.amname IN ('hnsw', 'ivfflat')")) { + statement.setString(1, this.schema); + statement.setString(2, table); + statement.setString(3, this.vectorField); + try (ResultSet rows = statement.executeQuery()) { + while (rows.next()) { + vectorIndexes.add(rows.getString(1) + " (" + rows.getString(2) + ")"); + usable |= metric.getOperatorClass().equals(rows.getString(2)); + } + } + } + if (!usable) { + mismatches.add( + "no " + + metric + + " (" + + metric.getOperatorClass() + + ") index on " + + quote(this.vectorField) + + (vectorIndexes.isEmpty() + ? "" + : "; existing vector indexes: " + vectorIndexes) + + "; similarity searches with this metric will scan the table"); + } + return mismatches; + } + + /** Drops the table for the given collection if it exists. */ + @Override + public synchronized void deleteCollection(String name) throws Exception { + String table = identifier("collection", name); + this.verifiedTables.removeIf(key -> key.startsWith(table + "|")); + try (Statement statement = connection().createStatement()) { + statement.execute("DROP TABLE IF EXISTS " + qualified(table)); + } catch (SQLException e) { + discardIfBroken(e); + throw e; + } + } + + /** + * Retrieve documents from the vector store. + * + * <p>If {@code ids} are provided, this method selects by primary key and ignores {@code + * filters} and {@code limit} per the {@link BaseVectorStore} contract; an empty id list returns + * no documents rather than scanning the table. Otherwise it selects rows whose metadata + * contains {@code filters}, up to {@code limit} rows (all matching rows when {@code limit} is + * null), ordered by id for a stable result. + */ + @Override + public synchronized List<Document> get( + @Nullable List<String> ids, + @Nullable String collection, + @Nullable Map<String, Object> filters, + @Nullable Integer limit, + Map<String, Object> extraArgs) + throws IOException { + String table = resolveCollection(collection); + if (ids != null && ids.isEmpty()) { + return Collections.emptyList(); + } + String select = + "SELECT " + + quote(this.idField) + + ", " + + quote(this.contentField) + + ", " + + quote(this.metadataField) + + " FROM " + + qualified(table); + try { + Connection connection = connection(); + if (ids != null) { + try (PreparedStatement statement = + connection.prepareStatement( + select + + " WHERE " + + quote(this.idField) + + " = ANY(?) ORDER BY " + + quote(this.idField))) { + statement.setArray(1, textArray(connection, ids)); + return readDocuments(statement, false); + } + } + String filterJson = filtersToJson(filters); + String sql = + select + + (filterJson == null + ? "" + : " WHERE " + quote(this.metadataField) + " @> ?::jsonb") + + " ORDER BY " + + quote(this.idField) + + (limit == null ? "" : " LIMIT ?"); + try (PreparedStatement statement = connection.prepareStatement(sql)) { + int parameter = 1; + if (filterJson != null) { + statement.setString(parameter++, filterJson); + } + if (limit != null) { + statement.setInt(parameter, positive("limit", limit)); + } + return readDocuments(statement, false); + } + } catch (SQLException e) { + discardIfBroken(e); + throw new IOException("Failed to read documents from " + qualified(table) + ".", e); + } + } + + /** + * Delete documents in the vector store. + * + * <p>If ids are provided, the corresponding rows are deleted; an empty id list is a no-op + * rather than a request to delete everything. Otherwise rows whose metadata contains {@code + * filters} are deleted; without filters every row in the table is deleted. + */ + @Override + public synchronized void delete( + @Nullable List<String> ids, + @Nullable String collection, + @Nullable Map<String, Object> filters, + Map<String, Object> extraArgs) + throws IOException { + String table = resolveCollection(collection); + if (ids != null && ids.isEmpty()) { + return; + } + try { + Connection connection = connection(); + if (ids != null) { + try (PreparedStatement statement = + connection.prepareStatement( + "DELETE FROM " + + qualified(table) + + " WHERE " + + quote(this.idField) + + " = ANY(?)")) { + statement.setArray(1, textArray(connection, ids)); + statement.executeUpdate(); + } + return; + } + String filterJson = filtersToJson(filters); + if (filterJson == null) { + try (Statement statement = connection.createStatement()) { + statement.executeUpdate("DELETE FROM " + qualified(table)); + } + return; + } + try (PreparedStatement statement = + connection.prepareStatement( + "DELETE FROM " + + qualified(table) + + " WHERE " + + quote(this.metadataField) + + " @> ?::jsonb")) { + statement.setString(1, filterJson); + statement.executeUpdate(); + } + } catch (SQLException e) { + discardIfBroken(e); + throw new IOException("Failed to delete documents from " + qualified(table) + ".", e); + } + } + + /** + * Executes a similarity search using a pre-computed embedding. + * + * <p>Rows are ordered by the distance operator of the metric (nearest first) and the metric's + * score is set on each {@link Document}. {@code metric_type} in {@code args} overrides the + * descriptor metric for this query; the vector index is only used when it matches. On pgvector + * 0.8+ the search runs with the configured {@code iterative_scan} mode (overridable in {@code + * args}) so that an index scan keeps going until {@code limit} rows match, filtered or not. + */ + @Override + public synchronized List<Document> queryEmbedding( + float[] embedding, + int limit, + @Nullable String collection, + @Nullable Map<String, Object> filters, + Map<String, Object> args) { + String table = resolveCollection(collection); + MetricType metric = enumFromMap(MetricType.class, args, "metric_type", this.metricType); + IterativeScan scan = + enumFromMap(IterativeScan.class, args, "iterative_scan", this.iterativeScan); + String filterJson = filtersToJson(filters); + String distance = quote(this.vectorField) + " " + metric.getOperator() + " ?::vector"; + // The distance column is referenced by position: any alias could clash with a user + // column. A relaxed iterative scan may emit hits slightly out of order, so the LIMIT runs + // in a materialized CTE and the outer query re-sorts the few rows it kept. + String sql = + "WITH hits AS MATERIALIZED (SELECT " + + quote(this.idField) + + ", " + + quote(this.contentField) + + ", " + + quote(this.metadataField) + + ", " + + distance + + " FROM " + + qualified(table) + + (filterJson == null + ? "" + : " WHERE " + quote(this.metadataField) + " @> ?::jsonb") + + " ORDER BY 4 LIMIT ?) SELECT * FROM hits ORDER BY 4"; Review Comment: With the default `RELAXED_ORDER`, PostgreSQL 17+ can omit the outer sort because it inherits the CTE’s ordering, leaving results out of distance order. Could we name the distance column and use `ORDER BY distance + 0`, as [pgvector recommends](https://github.com/pgvector/pgvector#iterative-index-scans)? I reproduced this on PG17; adding `+ 0` restored the expected order. ########## integrations/vector-stores/pgvector/src/test/java/org/apache/flink/agents/integrations/vectorstores/pgvector/PgVectorVectorStoreTest.java: ########## @@ -0,0 +1,1110 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.integrations.vectorstores.pgvector; + +import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup; +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.api.vectorstores.BaseVectorStore; +import org.apache.flink.agents.api.vectorstores.CollectionManageableVectorStore; +import org.apache.flink.agents.api.vectorstores.Document; +import org.apache.flink.agents.api.vectorstores.VectorStoreQuery; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.mockito.Mockito; + +import java.io.IOException; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.assertj.core.api.Assertions.within; + +/** + * Tests for {@link PgVectorVectorStore}. The live tests need a PostgreSQL server with the pgvector + * extension and are enabled by {@code PGVECTOR_URI} (a JDBC URL); {@code PGVECTOR_USERNAME} and + * {@code PGVECTOR_PASSWORD} default to {@code postgres}. {@code tools/docker/pgvector} starts one. + */ +public class PgVectorVectorStoreTest { + + private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null); + private static final String LIVE = "PGVECTOR_URI"; + + @Test + void testConstructorAndStoreKwargs() { + ResourceDescriptor desc = + ResourceDescriptor.Builder.newBuilder(PgVectorVectorStore.class.getName()) + .addInitialArgument("embedding_model", "embeddingModel") + .addInitialArgument("uri", "jdbc:postgresql://db:5432/agents") + .addInitialArgument("username", "user") + .addInitialArgument("password", "secret") + .addInitialArgument("schema", "rag") + .addInitialArgument("collection", "test_collection") + .addInitialArgument("dims", 5) + .addInitialArgument("metric_type", "ip") + .addInitialArgument("index_type", "ivfflat") + .addInitialArgument("index_params", Map.of("lists", 100)) + .addInitialArgument("create_extension", false) + .addInitialArgument("iterative_scan", "strict_order") + .build(); + PgVectorVectorStore store = new PgVectorVectorStore(desc, NOOP); + Map<String, Object> kwargs = store.getStoreKwargs(); + assertThat(store).isInstanceOf(BaseVectorStore.class); + assertThat(store).isInstanceOf(CollectionManageableVectorStore.class); + assertThat(kwargs) + .containsEntry("uri", "jdbc:postgresql://db:5432/agents") + .containsEntry("schema", "rag") + .containsEntry("collection", "test_collection") + .containsEntry("index", "test_collection") + .containsEntry("dims", 5) + .containsEntry("metric_type", "IP") + .containsEntry("index_type", "IVFFLAT") + .containsEntry("index_params", Map.of("lists", 100L)) + .containsEntry("create_extension", false) + .containsEntry("iterative_scan", "STRICT_ORDER"); + // Credentials never travel with the store kwargs. + assertThat(kwargs).doesNotContainKeys("username", "password"); + store.close(); + // A closed store does not quietly reopen a connection. + assertThatThrownBy(() -> store.get(List.of("x"), null, null, null, Map.of())) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("closed"); + } + + @Test + void testDefaultsAndAliases() { + PgVectorVectorStore store = + new PgVectorVectorStore( + ResourceDescriptor.Builder.newBuilder(PgVectorVectorStore.class.getName()) + .addInitialArgument("index", "aliased") + .build(), + NOOP); + assertThat(store.getStoreKwargs()) + .containsEntry("uri", PgVectorVectorStore.DEFAULT_URI) + .containsEntry("schema", PgVectorVectorStore.DEFAULT_SCHEMA) + .containsEntry("collection", "aliased") + .containsEntry("dims", PgVectorVectorStore.DEFAULT_DIMENSION) + .containsEntry("metric_type", "COSINE") + .containsEntry("index_type", "HNSW") + .containsEntry("index_params", Map.of()) + .containsEntry("create_extension", true) + .containsEntry("iterative_scan", "RELAXED_ORDER"); + + PgVectorVectorStore hostStore = + new PgVectorVectorStore( + ResourceDescriptor.Builder.newBuilder(PgVectorVectorStore.class.getName()) + .addInitialArgument("host", "db.internal") + .addInitialArgument("port", 6543) + .addInitialArgument("database", "agents") + .build(), + NOOP); + assertThat(hostStore.getStoreKwargs()) + .containsEntry("uri", "jdbc:postgresql://db.internal:6543/agents") + .containsEntry("collection", PgVectorVectorStore.DEFAULT_COLLECTION); + + // port/database are honoured without an explicit host. + PgVectorVectorStore portStore = + new PgVectorVectorStore( + ResourceDescriptor.Builder.newBuilder(PgVectorVectorStore.class.getName()) + .addInitialArgument("port", 6543) + .addInitialArgument("database", "agents") + .build(), + NOOP); + assertThat(portStore.getStoreKwargs()) + .containsEntry("uri", "jdbc:postgresql://localhost:6543/agents"); + + // Credentials embedded in the JDBC URL never surface through the store kwargs. + PgVectorVectorStore credentialStore = + newStore( + "uri", + "jdbc:postgresql://db/agents?user=svc&password=s3cret&ssl=true" + + "&sslkey=k.pk8&sslpassword=hunter2"); + assertThat(credentialStore.getStoreKwargs()) + .containsEntry("uri", "jdbc:postgresql://db/agents?ssl=true&sslkey=k.pk8") + .doesNotContainKeys("username", "password"); + assertThat(credentialStore.getStoreKwargs().toString()) + .doesNotContain("s3cret") + .doesNotContain("hunter2"); + assertThat(PgVectorVectorStore.stripCredentials("jdbc:postgresql://db/agents?password=x")) + .isEqualTo("jdbc:postgresql://db/agents"); + } + + @Test + void testArgumentValidation() { + assertThatThrownBy(() -> newStore("collection", "docs; DROP TABLE x")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("collection"); + assertThatThrownBy(() -> newStore("vector_field", "embedding\"")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("vector_field"); + assertThatThrownBy(() -> newStore("dims", 0)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dims"); + assertThatThrownBy(() -> newStore("dims", "abc")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dims"); + assertThatThrownBy(() -> newStore("dims", 768.5)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dims"); + assertThatThrownBy(() -> newStore("dims", 3_000_000_000L)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("dims"); + assertThat(newStore("dims", 768.0).getStoreKwargs()).containsEntry("dims", 768); + assertThatThrownBy(() -> newStore("metric_type", "DOT")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("metric_type"); + assertThatThrownBy(() -> newStore("index_type", "FLAT")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("index_type"); + assertThatThrownBy(() -> newStore("create_extension", "yes")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("create_extension"); + assertThat(newStore("create_extension", " FALSE ").getStoreKwargs()) + .containsEntry("create_extension", false); + assertThatThrownBy(() -> newStore("index_params", Map.of("m", "16.5"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("index_params"); + assertThatThrownBy(() -> newStore("index_params", Map.of("m", "16); DROP TABLE x"))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("index_params"); + assertThatThrownBy(() -> newStore("index_params", "m=16")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("index_params"); + assertThatThrownBy(() -> newStore("uri", "postgresql://svc:S3cret@localhost/postgres")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("jdbc:postgresql:") + .hasMessageContaining("//***@localhost") + .satisfies(e -> assertThat(e.getMessage()).doesNotContain("S3cret")); + assertThat(PgVectorVectorStore.redactUri("jdbc:postgresql://db/x?password=p&ssl=true")) + .isEqualTo("jdbc:postgresql://db/x?ssl=true"); + // pgjdbc's short and socket forms are accepted as given. + assertThat(newStore("uri", "jdbc:postgresql:agents").getStoreKwargs()) + .containsEntry("uri", "jdbc:postgresql:agents"); + assertThat( + newStore("uri", "jdbc:postgresql:///agents?host=/var/run/postgresql") + .getStoreKwargs()) + .containsEntry("uri", "jdbc:postgresql:///agents?host=/var/run/postgresql"); + assertThatThrownBy(() -> newStore("collection", "t".repeat(64))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("63 characters"); + // Wide vectors are only rejected when a table with an index would actually be created + // (checked live below); the constructor and the schema check accept them. + assertThat(newStore("dims", 3072).getStoreKwargs()).containsEntry("dims", 3072); + assertThatThrownBy(() -> newStore("connect_timeout_s", -1)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("connect_timeout_s"); + assertThatThrownBy(() -> newStore("socket_timeout_s", "soon")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("socket_timeout_s"); + assertThatThrownBy(() -> newStore("index_params", Map.of("m", 16.5))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("must be an integer"); + assertThatThrownBy(() -> newStore("iterative_scan", "fast")) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("iterative_scan"); + assertThat(PgVectorVectorStore.isAtLeast("0.8.6", new int[] {0, 8})).isTrue(); + assertThat(PgVectorVectorStore.isAtLeast("v0.8.1", new int[] {0, 8})).isTrue(); + assertThat(PgVectorVectorStore.isAtLeast("0.8", new int[] {0, 8})).isTrue(); + assertThat(PgVectorVectorStore.isAtLeast("1.0.0", new int[] {0, 8})).isTrue(); + assertThat(PgVectorVectorStore.isAtLeast("0.7.4", new int[] {0, 8})).isFalse(); + assertThat(PgVectorVectorStore.isAtLeast(null, new int[] {0, 8})).isFalse(); + assertThatThrownBy( + () -> + new PgVectorVectorStore( + ResourceDescriptor.Builder.newBuilder( + PgVectorVectorStore.class.getName()) + .addInitialArgument("uri", "jdbc:postgresql://a/db") + .addInitialArgument("host", "b") + .build(), + NOOP)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("uri and host"); + assertThatThrownBy(() -> newStore("port", 0)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("port"); + // Indexed columns must leave room for a unique index name on any table. + int maxColumn = PgVectorVectorStore.MAX_INDEXED_COLUMN_LENGTH; + assertThatThrownBy(() -> newStore("vector_field", "v".repeat(maxColumn + 1))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("index name"); + assertThat( + PgVectorVectorStore.indexName( + "t".repeat(63), + newStore("vector_field", "v".repeat(maxColumn)) + .getStoreKwargs() + .get("vector_field") + .toString()) + .length()) + .isLessThanOrEqualTo(63); + // Wide vectors are fine without an index; integral floats normalize to integers. + // Timeouts are connection settings and stay out of the store kwargs. + PgVectorVectorStore wide = + new PgVectorVectorStore( + ResourceDescriptor.Builder.newBuilder(PgVectorVectorStore.class.getName()) + .addInitialArgument("dims", 3072) + .addInitialArgument("index_type", "NONE") + .addInitialArgument("index_params", Map.of("m", 16.0)) + .addInitialArgument("connect_timeout_s", 3) + .addInitialArgument("socket_timeout_s", 30) + .build(), + NOOP); + assertThat(wide.getStoreKwargs()) + .containsEntry("index_params", Map.of("m", 16L)) + .doesNotContainKeys("connect_timeout_s", "socket_timeout_s"); + } + + @Test + void testIndexNamesStayUniqueWithinIdentifierLimit() { + assertThat(PgVectorVectorStore.indexName("docs", "embedding")) + .isEqualTo("docs_embedding_idx"); + String longTable = "fa_pgvector_" + "x".repeat(50); + String vectorIndex = PgVectorVectorStore.indexName(longTable, "embedding"); + String metadataIndex = PgVectorVectorStore.indexName(longTable, "metadata"); + assertThat(vectorIndex.length()).isLessThanOrEqualTo(63); + assertThat(metadataIndex.length()).isLessThanOrEqualTo(63); + assertThat(vectorIndex).endsWith("_embedding_idx").isNotEqualTo(metadataIndex); + // Distinct long tables never share an index name. + assertThat(PgVectorVectorStore.indexName(longTable + "a", "embedding")) + .isNotEqualTo(PgVectorVectorStore.indexName(longTable + "b", "embedding")); + } + + @Test + void testFiltersCompileToContainmentJson() { + PgVectorVectorStore store = newStore("collection", "docs"); + assertThat(store.filtersToJson(null)).isNull(); + assertThat(store.filtersToJson(Map.of())).isNull(); + assertThat(store.filtersToJson(Map.of("category", "docs"))) + .isEqualTo("{\"category\":\"docs\"}"); + assertThat(store.filtersToJson(Map.of("count", 3))).isEqualTo("{\"count\":3}"); + assertThatThrownBy(() -> store.filtersToJson(Map.of("nested", Map.of("a", 1)))) + .isInstanceOf(UnsupportedOperationException.class); + // A list would turn containment into a subset match, which is not equality. + assertThatThrownBy(() -> store.filtersToJson(Map.of("tags", List.of("a")))) + .isInstanceOf(UnsupportedOperationException.class); + Map<String, Object> withNull = new HashMap<>(); + withNull.put("category", null); + assertThatThrownBy(() -> store.filtersToJson(withNull)) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + void testVectorLiteralAndScores() { + assertThat(PgVectorVectorStore.vectorLiteral(new float[] {1.0f, -0.5f, 0.25f})) + .isEqualTo("[1.0,-0.5,0.25]"); + assertThat(PgVectorVectorStore.vectorLiteral(new float[0])).isEqualTo("[]"); + assertThatThrownBy(() -> PgVectorVectorStore.vectorLiteral(new float[] {Float.NaN})) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("finite"); + // Cosine distance 0.25 -> similarity 0.75; L2 keeps the distance; IP negates <#>. + assertThat(PgVectorVectorStore.MetricType.COSINE.toScore(0.25)) + .isCloseTo(0.75f, within(1e-6f)); + assertThat(PgVectorVectorStore.MetricType.L2.toScore(2.5)).isCloseTo(2.5f, within(1e-6f)); + assertThat(PgVectorVectorStore.MetricType.IP.toScore(-3.0)).isCloseTo(3.0f, within(1e-6f)); + assertThat(PgVectorVectorStore.MetricType.COSINE.getOperator()).isEqualTo("<=>"); + assertThat(PgVectorVectorStore.MetricType.L2.getOperatorClass()).isEqualTo("vector_l2_ops"); + } + + @Test + @EnabledIfEnvironmentVariable(named = LIVE, matches = ".+") + void testCollectionManagement() throws Exception { + String collection = collectionName("collection_management"); + PgVectorVectorStore store = openStore(collection); + try { + store.createCollectionIfNotExists(collection, Map.of()); + // Creating again is a no-op, with or without options. + store.createCollectionIfNotExists(collection, Map.of()); + store.createCollectionIfNotExists(collection, null); + assertThat(store.schemaMismatches(collection, null)).isEmpty(); + Assertions.assertTrue( + store.get(null, collection, null, 10, Collections.emptyMap()).isEmpty()); + assertThat(indexes(collection)) + .containsEntry( + PgVectorVectorStore.indexName(collection, "embedding"), + "hnsw/vector_cosine_ops") + .containsEntry( + PgVectorVectorStore.indexName(collection, "metadata"), "gin/jsonb_ops"); + store.deleteCollection(collection); + Assertions.assertThrows( + Exception.class, + () -> store.get(null, collection, null, 10, Collections.emptyMap())); + } finally { + dropCollectionQuietly(store, collection); + store.close(); + } + } + + @Test + @EnabledIfEnvironmentVariable(named = LIVE, matches = ".+") + void testExistingTableMismatchesAreReported() throws Exception { + String collection = collectionName("mismatch"); + PgVectorVectorStore store = openStore(collection); + try { + store.createCollectionIfNotExists(collection, Map.of("metric_type", "L2")); + assertThat(store.schemaMismatches(collection, Map.of("metric_type", "L2"))).isEmpty(); + // IF NOT EXISTS keeps the L2 table; a COSINE store on it has no usable index and a + // different width, and both facts are reported (and logged by create). + store.createCollectionIfNotExists(collection, Map.of("dims", 7)); + // ... and nothing was added to it: still one vector index, the L2 one. + assertThat(indexes(collection).values()).containsOnlyOnce("hnsw/vector_l2_ops"); + assertThat(indexes(collection).values()).doesNotContain("hnsw/vector_cosine_ops"); + List<String> mismatches = store.schemaMismatches(collection, Map.of("dims", 7)); + assertThat(mismatches).hasSize(2); + assertThat(mismatches.get(0)).contains("vector(5)").contains("vector(7)"); + assertThat(mismatches.get(1)).contains("vector_cosine_ops").contains("vector_l2_ops"); + // A store configured for wider vectors than an index allows can still describe and + // "ensure" the existing table; only creating a new table enforces the index limit. + PgVectorVectorStore wide = openStore(collection); + try { + wide.createCollectionIfNotExists(collection, Map.of("dims", 3072)); + assertThat(wide.schemaMismatches(collection, Map.of("dims", 3072))) + .isNotEmpty() + .noneMatch(m -> m.contains("2000-dimension")); + assertThatThrownBy( + () -> + wide.createCollectionIfNotExists( + collectionName("wide"), Map.of("dims", 3072))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("2000-dimension limit"); + } finally { + wide.close(); + } + // Without an index requirement only the column is checked. + assertThat(store.schemaMismatches(collection, Map.of("index_type", "NONE"))).isEmpty(); + // A non-table relation squatting on the name is an error, not "create it". + String viewName = collectionName("view"); + try (Connection admin = adminConnection(); + Statement statement = admin.createStatement()) { + statement.execute("CREATE VIEW public.\"" + viewName + "\" AS SELECT 1 AS id"); + } + try { + assertThatThrownBy(() -> store.createCollectionIfNotExists(viewName, Map.of())) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("not a table"); + } finally { + try (Connection admin = adminConnection(); + Statement statement = admin.createStatement()) { + statement.execute("DROP VIEW IF EXISTS public.\"" + viewName + "\""); + } + } + // Unknown collections are reported rather than assumed fine. + assertThat(store.schemaMismatches(collectionName("absent"), Map.of())) + .singleElement() + .asString() + .contains("does not exist"); + } finally { + dropCollectionQuietly(store, collection); + store.close(); + } + } + + @Test + @EnabledIfEnvironmentVariable(named = LIVE, matches = ".+") + void testCreateCollectionHonoursIndexArguments() throws Exception { + String collection = collectionName("index_args"); + PgVectorVectorStore store = openStore(collection); + try { + store.createCollectionIfNotExists( + collection, + Map.of( + "metric_type", "L2", + "index_type", "IVFFLAT", + "index_params", Map.of("lists", 4))); + assertThat(indexes(collection)) + .containsEntry( + PgVectorVectorStore.indexName(collection, "embedding"), + "ivfflat/vector_l2_ops"); + try (Connection connection = adminConnection(); + Statement statement = connection.createStatement(); + ResultSet rows = + statement.executeQuery( + "SELECT indexdef FROM pg_indexes WHERE indexname = '" + + PgVectorVectorStore.indexName(collection, "embedding") + + "'")) { + Assertions.assertTrue(rows.next(), "vector index should exist"); + assertThat(rows.getString(1)).contains("lists='4'"); + } + } finally { + dropCollectionQuietly(store, collection); + store.close(); + } + } + + @Test + @EnabledIfEnvironmentVariable(named = LIVE, matches = ".+") + void testDocumentManagement() throws Exception { + String collection = collectionName("document_management"); + PgVectorVectorStore store = openStore(collection); + try { + store.createCollectionIfNotExists(collection, Map.of()); + store.add( + List.of( + new Document( + "pgvector is a PostgreSQL extension", + Map.of("category", "database", "source", "test"), + "doc1"), + new Document( + "Apache Flink Agents is an AI framework", + Map.of("category", "ai-agent", "source", "test"), + "doc2")), + collection, + Collections.emptyMap()); + + List<Document> all = store.get(null, collection, null, 10, Collections.emptyMap()); + Assertions.assertEquals(2, all.size()); + assertDocument( + documentById(all, "doc1"), + "doc1", + "pgvector is a PostgreSQL extension", + Map.of("category", "database", "source", "test")); + assertDocument( + documentById(all, "doc2"), + "doc2", + "Apache Flink Agents is an AI framework", + Map.of("category", "ai-agent", "source", "test")); + + List<Document> byId = + store.get(List.of("doc1"), collection, null, null, Collections.emptyMap()); + Assertions.assertEquals(1, byId.size()); + Assertions.assertEquals("doc1", byId.get(0).getId()); + // An empty id list returns nothing instead of scanning the table. + Assertions.assertTrue( + store.get(List.of(), collection, null, null, Collections.emptyMap()).isEmpty()); + + Assertions.assertEquals( + 1, store.get(null, collection, null, 1, Collections.emptyMap()).size()); + + // An empty id list is a no-op, not a request to delete everything. + store.delete(List.of(), collection, null, Collections.emptyMap()); + Assertions.assertEquals( + 2, store.get(null, collection, null, 10, Collections.emptyMap()).size()); + + store.delete(List.of("doc1"), collection, null, Collections.emptyMap()); + List<Document> remaining = + store.get(null, collection, null, 10, Collections.emptyMap()); + Assertions.assertEquals(1, remaining.size()); + Assertions.assertEquals("doc2", remaining.get(0).getId()); + + store.delete(null, collection, null, Collections.emptyMap()); + Assertions.assertTrue( + store.get(null, collection, null, 10, Collections.emptyMap()).isEmpty()); + } finally { + dropCollectionQuietly(store, collection); + store.close(); + } + } + + @Test + @EnabledIfEnvironmentVariable(named = LIVE, matches = ".+") + void testFiltersDsl() throws Exception { + String collection = collectionName("filters_dsl"); + PgVectorVectorStore store = openStore(collection); + try { + store.createCollectionIfNotExists(collection, Map.of()); + store.add( + List.of( + new Document( + "pgvector is a PostgreSQL extension", + Map.of("category", "database", "user_id", "alice", "rank", 1), + "doc_alice"), + new Document( + "Apache Flink Agents is an AI framework", + Map.of("category", "ai-agent", "user_id", "bob", "rank", 2), + "doc_bob")), + collection, + Collections.emptyMap()); + + List<Document> aliceOnly = + store.get(null, collection, Map.of("user_id", "alice"), 10, Map.of()); + Assertions.assertEquals(1, aliceOnly.size()); + Assertions.assertEquals("doc_alice", aliceOnly.get(0).getId()); + + // Numeric equality and two-key conjunction. + Assertions.assertEquals( + "doc_bob", + store.get(null, collection, Map.of("rank", 2), 10, Map.of()).get(0).getId()); + Assertions.assertTrue( + store.get( + null, + collection, + Map.of("user_id", "alice", "category", "ai-agent"), + 10, + Map.of()) + .isEmpty()); + + // A filter applies to similarity search too: the closest vector belongs to bob, but + // only alice's document may be returned. + List<Document> aliceQueried = + store.queryEmbedding( + new float[] {0.0f, 1.0f, 0.0f, 0.0f, 0.0f}, + 5, + collection, + Map.of("user_id", "alice"), + Collections.emptyMap()); + Assertions.assertEquals(1, aliceQueried.size()); + Assertions.assertEquals("doc_alice", aliceQueried.get(0).getId()); + + store.delete(null, collection, Map.of("user_id", "alice"), Collections.emptyMap()); + List<Document> remaining = store.get(null, collection, null, 10, Map.of()); + Assertions.assertEquals(1, remaining.size()); + Assertions.assertEquals("doc_bob", remaining.get(0).getId()); + } finally { + dropCollectionQuietly(store, collection); + store.close(); + } + } + + @Test + @EnabledIfEnvironmentVariable(named = LIVE, matches = ".+") + void testAddGeneratesIdsAndUpdateOverwrites() throws Exception { + String collection = collectionName("add_update"); + PgVectorVectorStore store = openStore(collection); + try { + store.createCollectionIfNotExists(collection, Map.of()); + // add treats "" as "generate an id"; update must not upsert a row keyed "". + assertThatThrownBy( + () -> + store.update( + List.of(new Document("x", Map.of(), "")), + collection, + Collections.emptyMap())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("id"); + List<String> ids = + store.add( + List.of( + new Document( + "pgvector is a PostgreSQL extension", + Map.of("category", "database"), + null)), + collection, + Collections.emptyMap()); + Assertions.assertEquals(1, ids.size()); + Assertions.assertFalse(ids.get(0).isEmpty()); + List<Document> stored = + store.get(List.of(ids.get(0)), collection, null, null, Collections.emptyMap()); + Assertions.assertEquals(1, stored.size()); + + Document original = + new Document( + "pgvector is a PostgreSQL extension", + Map.of("category", "database"), + "doc1"); + store.add(List.of(original), collection, Collections.emptyMap()); + // Adding again with an existing id replaces the row instead of failing, so a replayed + // batch is idempotent; the new document in the same batch is stored as well. + Document fresh = + new Document( + "Apache Flink Agents is an AI framework", + Map.of("category", "ai"), + "fresh"); + Document replayed = + new Document( + "pgvector is a PostgreSQL extension", + Map.of("category", "database", "replayed", true), + "doc1"); + Assertions.assertEquals( + List.of("fresh", "doc1"), + store.add(List.of(fresh, replayed), collection, Collections.emptyMap())); + // The generated-id document from above, doc1 (replaced in place) and fresh. + Assertions.assertEquals( + 3, store.get(null, collection, null, null, Collections.emptyMap()).size()); + Assertions.assertEquals( + true, + store.get(List.of("doc1"), collection, null, null, Collections.emptyMap()) + .get(0) + .getMetadata() + .get("replayed")); + // The same id twice in one batch stores one row, the later document winning. + store.add( + List.of( + new Document("first", Map.of(), "twice", new float[] {0, 0, 0, 1, 0}), + new Document("second", Map.of(), "twice", new float[] {0, 0, 0, 0, 1})), + collection, + Collections.emptyMap()); + List<Document> twice = + store.get(List.of("twice"), collection, null, null, Collections.emptyMap()); + Assertions.assertEquals(1, twice.size()); + Assertions.assertEquals("second", twice.get(0).getContent()); + // A batch is still atomic: an unserialisable document fails the whole batch. + Map<String, Object> unserialisable = new HashMap<>(); + unserialisable.put("self", new Object()); + assertThatThrownBy( + () -> + store.add( + List.of( + new Document( + "a", + Map.of(), + "atomic_a", + new float[] {1, 0, 0, 0, 0}), + new Document( + "b", + unserialisable, + "atomic_b", + new float[] {0, 1, 0, 0, 0})), + collection, + Collections.emptyMap())) + .isInstanceOf(IOException.class) + .hasMessageContaining("add"); + Assertions.assertTrue( + store.get(List.of("atomic_a"), collection, null, null, Collections.emptyMap()) + .isEmpty()); + + Document rewritten = + new Document( + "pgvector stores dense vectors", Map.of("category", "updated"), "doc1"); + store.update(List.of(rewritten), collection, Collections.emptyMap()); + List<Document> after = + store.get(List.of("doc1"), collection, null, null, Collections.emptyMap()); + Assertions.assertEquals(1, after.size()); + Assertions.assertEquals("pgvector stores dense vectors", after.get(0).getContent()); + Assertions.assertEquals("updated", after.get(0).getMetadata().get("category")); + // The embedding was replaced too: the rewritten text maps to the second axis. + List<Document> hits = + store.queryEmbedding( + new float[] {0.0f, 1.0f, 0.0f, 0.0f, 0.0f}, + 1, + collection, + null, + Collections.emptyMap()); + Assertions.assertEquals("doc1", hits.get(0).getId()); Review Comment: `fresh` and the updated `doc1` have identical embeddings, so either can be returned by this query with `limit=1`. This assertion failed on PostgreSQL 17 because it returned `fresh`. Could we give `fresh` a different embedding so that `doc1` is the unique nearest match? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
