From 6bd059a8b7f213c27a59dfbc5f24670f257148ee Mon Sep 17 00:00:00 2001 From: Muthukumaran Navaneethakrishnan Date: Thu, 8 Aug 2024 09:05:44 +0530 Subject: [PATCH] Fix IN/NOT IN filters for PgVector JSON queries PgVectorFilterExpressionConverter was generating incorrect SQL for IN and NOT IN filters with PostgreSQL JSON data types. This caused BadSqlGrammarException errors when executing queries. This change modifies the converter to generate correct SQL syntax for these operations, ensuring compatibility with PostgreSQL's JSON handling capabilities. Why: - Improves query reliability for PgVector stores - Enables more complex filtering operations on JSON data - Eliminates unexpected errors in query execution Fixes #1179 --- .../PgVectorFilterExpressionConverter.java | 50 ++++++++++++++++--- ...gVectorFilterExpressionConverterTests.java | 9 ++-- .../ai/vectorstore/PgVectorStoreIT.java | 44 ++++++++++++++++ 3 files changed, 93 insertions(+), 10 deletions(-) diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java index 0b2ccf2fd..18d8e23bc 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverter.java @@ -15,24 +15,64 @@ */ package org.springframework.ai.vectorstore; +import org.springframework.ai.vectorstore.filter.Filter; import org.springframework.ai.vectorstore.filter.Filter.Expression; import org.springframework.ai.vectorstore.filter.Filter.Group; import org.springframework.ai.vectorstore.filter.Filter.Key; import org.springframework.ai.vectorstore.filter.converter.AbstractFilterExpressionConverter; +import java.util.List; /** * Converts {@link Expression} into PgVector metadata filter expression format. * (https://www.postgresql.org/docs/current/functions-json.html) * + * @author Muthukumaran Navaneethakrishnan * @author Christian Tzolov */ public class PgVectorFilterExpressionConverter extends AbstractFilterExpressionConverter { @Override protected void doExpression(Expression expression, StringBuilder context) { - this.convertOperand(expression.left(), context); - context.append(getOperationSymbol(expression)); - this.convertOperand(expression.right(), context); + if (expression.type() == Filter.ExpressionType.IN) { + handleIn(expression, context); + } + else if (expression.type() == Filter.ExpressionType.NIN) { + handleNotIn(expression, context); + } + else { + this.convertOperand(expression.left(), context); + context.append(getOperationSymbol(expression)); + this.convertOperand(expression.right(), context); + } + } + + private void handleIn(Expression expression, StringBuilder context) { + context.append("("); + convertToConditions(expression, context); + context.append(")"); + } + + private void convertToConditions(Expression expression, StringBuilder context) { + Filter.Value right = (Filter.Value) expression.right(); + Object value = right.value(); + if (!(value instanceof List)) { + throw new IllegalArgumentException("Expected a List, but got: " + value.getClass().getSimpleName()); + } + List values = (List) value; + for (int i = 0; i < values.size(); i++) { + this.convertOperand(expression.left(), context); + context.append(" == "); + this.doSingleValue(values.get(i), context); + if (i < values.size() - 1) { + context.append(" || "); + } + } + } + + private void handleNotIn(Expression expression, StringBuilder context) { + context.append("!("); + convertToConditions(expression, context); + context.append(")"); } private String getOperationSymbol(Expression exp) { @@ -53,10 +93,6 @@ public class PgVectorFilterExpressionConverter extends AbstractFilterExpressionC return " > "; case GTE: return " >= "; - case IN: - return " in "; - case NIN: - return " nin "; default: throw new RuntimeException("Not supported expression type: " + exp.type()); } diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java index 9cdc6b7bd..ca662ab8e 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorFilterExpressionConverterTests.java @@ -34,6 +34,7 @@ import org.springframework.ai.vectorstore.filter.Filter.Value; import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; /** + * @author Muthukumaran Navaneethakrishnan * @author Christian Tzolov */ public class PgVectorFilterExpressionConverterTests { @@ -61,7 +62,8 @@ public class PgVectorFilterExpressionConverterTests { // genre in ["comedy", "documentary", "drama"] String vectorExpr = converter.convertExpression( new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama")))); - assertThat(vectorExpr).isEqualTo("$.genre in [\"comedy\",\"documentary\",\"drama\"]"); + assertThat(vectorExpr) + .isEqualTo("($.genre == \"comedy\" || $.genre == \"documentary\" || $.genre == \"drama\")"); } @Test @@ -82,7 +84,7 @@ public class PgVectorFilterExpressionConverterTests { new Expression(EQ, new Key("country"), new Value("BG")))), new Expression(NIN, new Key("city"), new Value(List.of("Sofia", "Plovdiv"))))); assertThat(vectorExpr) - .isEqualTo("($.year >= 2020 || $.country == \"BG\") && $.city nin [\"Sofia\",\"Plovdiv\"]"); + .isEqualTo("($.year >= 2020 || $.country == \"BG\") && !($.city == \"Sofia\" || $.city == \"Plovdiv\")"); } @Test @@ -93,7 +95,8 @@ public class PgVectorFilterExpressionConverterTests { new Expression(GTE, new Key("year"), new Value(2020))), new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US"))))); - assertThat(vectorExpr).isEqualTo("$.isOpen == true && $.year >= 2020 && $.country in [\"BG\",\"NL\",\"US\"]"); + assertThat(vectorExpr).isEqualTo( + "$.isOpen == true && $.year >= 2020 && ($.country == \"BG\" || $.country == \"NL\" || $.country == \"US\")"); } @Test diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java index 19586503d..cfdcf86d9 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/PgVectorStoreIT.java @@ -24,12 +24,15 @@ import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.UUID; +import java.util.stream.Stream; import javax.sql.DataSource; import org.junit.Assert; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; @@ -57,6 +60,7 @@ import org.testcontainers.junit.jupiter.Testcontainers; import com.zaxxer.hikari.HikariDataSource; /** + * @author Muthukumaran Navaneethakrishnan * @author Christian Tzolov */ @Testcontainers @@ -128,6 +132,46 @@ public class PgVectorStoreIT { }); } + static Stream provideFilters() { + return Stream.of(Arguments.of("country in ['BG','NL']", 3), // String Filters In + Arguments.of("year in [2020]", 1), // Numeric Filters In + Arguments.of("country not in ['BG']", 1), // String Filter Not In + Arguments.of("year not in [2020]", 2) // Numeric Filter Not In + ); + } + + @ParameterizedTest(name = "Filter expression {0} should return {1} records ") + @MethodSource("provideFilters") + public void searchWithInFilter(String expression, Integer expectedRecords) { + + contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=COSINE_DISTANCE") + .run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2020, "foo bar 1", "bar.foo")); + var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "NL")); + var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner", + Map.of("country", "BG", "year", 2023)); + + vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2)); + + SearchRequest searchRequest = SearchRequest.query("The World") + .withFilterExpression(expression) + .withTopK(5) + .withSimilarityThresholdAll(); + + List results = vectorStore.similaritySearch(searchRequest); + + assertThat(results).hasSize(expectedRecords); + + // Remove all documents from the store + dropTable(context); + }); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "COSINE_DISTANCE", "EUCLIDEAN_DISTANCE", "NEGATIVE_INNER_PRODUCT" }) public void searchWithFilters(String distanceType) {