Skip to content

Commit

Permalink
Add filter-based deletion to MariaDB vector store (#2125)
Browse files Browse the repository at this point in the history
Add string-based filter deletion alongside the Filter.Expression-based deletion
for MariaDB vector store, providing consistent deletion capabilities with
other vector store implementations.

Key changes:
- Add delete(Filter.Expression) implementation for MariaDB store
- Integrate with existing MariaDBFilterExpressionConverter
- Add comprehensive integration tests for filter deletion
- Support both simple and complex filter expressions

This maintains consistency with other vector store implementations and
enables flexible document deletion based on metadata filters.
  • Loading branch information
sobychacko authored Jan 28, 2025
1 parent c33edc7 commit fc1e90c
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import org.springframework.ai.util.JacksonUtils;
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionConverter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
Expand Down Expand Up @@ -133,6 +134,7 @@
*
* @author Diego Dupin
* @author Ilayaperumal Gopinathan
* @author Soby Chacko
* @since 1.0.0
*/
public class MariaDBVectorStore extends AbstractObservationVectorStore implements InitializingBean {
Expand Down Expand Up @@ -325,6 +327,25 @@ public Optional<Boolean> doDelete(List<String> idList) {
return Optional.of(updateCount == idList.size());
}

@Override
protected void doDelete(Filter.Expression filterExpression) {
Assert.notNull(filterExpression, "Filter expression must not be null");

try {
String nativeFilterExpression = this.filterExpressionConverter.convertExpression(filterExpression);

String sql = String.format("DELETE FROM %s WHERE %s", getFullyQualifiedTableName(), nativeFilterExpression);

logger.debug("Executing delete with filter: {}", sql);

this.jdbcTemplate.update(sql);
}
catch (Exception e) {
logger.error("Failed to delete documents by filter: {}", e.getMessage(), e);
throw new IllegalStateException("Failed to delete documents by filter", e);
}
}

@Override
public List<Document> doSimilaritySearch(SearchRequest request) {

Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright 2023-2024 the original author or authors.
* Copyright 2023-2025 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand All @@ -23,12 +23,14 @@
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import javax.sql.DataSource;

import com.zaxxer.hikari.HikariDataSource;
import org.junit.Assert;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
Expand All @@ -44,6 +46,7 @@
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser.FilterExpressionParseException;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
Expand All @@ -63,6 +66,7 @@

/**
* @author Diego Dupin
* @author Soby Chacko
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
Expand Down Expand Up @@ -357,6 +361,106 @@ public void searchWithThreshold(String distanceType) {
});
}

@Test
public void deleteByFilter() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").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));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "NL", "year", 2021));
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.builder()
.query("The World")
.topK(5)
.similarityThresholdAll()
.build();

List<Document> results = vectorStore.similaritySearch(searchRequest);
assertThat(results).hasSize(3);

Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ,
new Filter.Key("country"), new Filter.Value("BG"));

vectorStore.delete(filterExpression);

// Verify deletion - should only have NL document remaining
results = vectorStore.similaritySearch(searchRequest);
assertThat(results).hasSize(1);
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");

dropTable(context);
});
}

@Test
public void deleteWithStringFilterExpression() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").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));
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "NL", "year", 2021));
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));

var searchRequest = SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build();

List<Document> results = vectorStore.similaritySearch(searchRequest);
assertThat(results).hasSize(3);

vectorStore.delete("country == 'BG'");

results = vectorStore.similaritySearch(searchRequest);
assertThat(results).hasSize(1);
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");

dropTable(context);
});
}

@Test
public void deleteWithComplexFilterExpression() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);

var doc1 = new Document("Content 1", Map.of("type", "A", "priority", 1));
var doc2 = new Document("Content 2", Map.of("type", "A", "priority", 2));
var doc3 = new Document("Content 3", Map.of("type", "B", "priority", 1));

vectorStore.add(List.of(doc1, doc2, doc3));

// Complex filter expression: (type == 'A' AND priority > 1)
Filter.Expression priorityFilter = new Filter.Expression(Filter.ExpressionType.GT,
new Filter.Key("priority"), new Filter.Value(1));
Filter.Expression typeFilter = new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("type"),
new Filter.Value("A"));
Filter.Expression complexFilter = new Filter.Expression(Filter.ExpressionType.AND, typeFilter,
priorityFilter);

vectorStore.delete(complexFilter);

var results = vectorStore
.similaritySearch(SearchRequest.builder().query("Content").topK(5).similarityThresholdAll().build());

assertThat(results).hasSize(2);
assertThat(results.stream().map(doc -> doc.getMetadata().get("type")).collect(Collectors.toList()))
.containsExactlyInAnyOrder("A", "B");
assertThat(results.stream().map(doc -> doc.getMetadata().get("priority")).collect(Collectors.toList()))
.containsExactlyInAnyOrder(1, 1);

dropTable(context);
});
}

@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {
Expand Down

0 comments on commit fc1e90c

Please sign in to comment.