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
This commit is contained in:
Muthukumaran Navaneethakrishnan
2024-08-08 09:05:44 +05:30
committed by Mark Pollack
parent bc3f9acb86
commit 6bd059a8b7
3 changed files with 93 additions and 10 deletions

View File

@@ -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<Object> 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());
}

View File

@@ -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

View File

@@ -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<Arguments> 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<Document> 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) {