Add filter-based deletion to Cassandra Vector Store (#2143)

Implement filter-based deletion for CassandraVectorStore to support more
    flexible document removal based on metadata filters.

    Key changes:

    - Add doDelete(Filter.Expression) implementation
    - Implement workaround for Cassandra's limitations in direct filtered deletion
    - Fetch matching documents first, then delete by IDs
    - Add integration tests for various filter deletion scenarios
    - Support simple and complex filter expressions
    - Includes tests for:
      * Deleting by simple equality filter
      * Deleting by string-based filter expression
      * Deleting by complex AND filter expression

    Provides consistent deletion capabilities across vector store implementations
    while addressing Cassandra-specific constraints.

Signed-off-by: Soby Chacko <soby.chacko@broadcom.com>
This commit is contained in:
Soby Chacko
2025-01-29 15:42:55 -05:00
committed by GitHub
parent bd52786de7
commit b9549d5ed2
2 changed files with 148 additions and 1 deletions

View File

@@ -31,6 +31,7 @@ import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import com.datastax.oss.driver.api.core.CqlSession;
@@ -73,6 +74,7 @@ import org.springframework.ai.observation.conventions.VectorStoreSimilarityMetri
import org.springframework.ai.vectorstore.AbstractVectorStoreBuilder;
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.FilterExpressionConverter;
import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
@@ -315,6 +317,44 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
return Optional.of(Boolean.TRUE);
}
@Override
protected void doDelete(Filter.Expression filterExpression) {
Assert.notNull(filterExpression, "Filter expression must not be null");
try {
// TODO - Investigate why we can't do a direct filter based delete in
// Cassandra
// This SO thread seems to indicate that this is not possible in Cassandra
// https://stackoverflow.com/questions/70953262/unable-to-delete-multiple-rows-getting-some-partition-key-parts-are-missing-i
// Needs more research into this matter.
SearchRequest searchRequest = SearchRequest.builder()
.query("") // empty query since we only want filter matches
.filterExpression(filterExpression)
.topK(1000) // large enough to get all matches
.similarityThresholdAll()
.build();
List<Document> matchingDocs = similaritySearch(searchRequest);
if (!matchingDocs.isEmpty()) {
// Then delete those documents by ID
List<String> idsToDelete = matchingDocs.stream().map(Document::getId).collect(Collectors.toList());
Optional<Boolean> result = delete(idsToDelete);
if (result.isPresent() && !result.get()) {
throw new IllegalStateException("Failed to delete some documents");
}
logger.debug(() -> "Deleted " + idsToDelete.size() + " documents matching filter expression");
}
}
catch (Exception e) {
logger.error(e, () -> "Failed to delete documents by filter");
throw new IllegalStateException("Failed to delete documents by filter", e);
}
}
@Override
public List<Document> doSimilaritySearch(SearchRequest request) {
Preconditions.checkArgument(request.getTopK() <= 1000);

View File

@@ -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.
@@ -22,6 +22,7 @@ import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
@@ -42,6 +43,7 @@ import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumn;
import org.springframework.ai.vectorstore.cassandra.CassandraVectorStore.SchemaColumnTags;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -57,6 +59,7 @@ import static org.assertj.core.api.Assertions.assertThat;
*
* @author Mick Semb Wever
* @author Thomas Vitale
* @author Soby Chacko
* @since 1.0.0
*/
@Testcontainers
@@ -417,6 +420,110 @@ class CassandraVectorStoreIT {
});
}
@Test
void deleteByFilter() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createTestStore(context,
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),
new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) {
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", (short) 2020));
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", (short) 2023));
store.add(List.of(bgDocument, nlDocument, bgDocument2));
// Verify initial state
List<Document> results = store
.similaritySearch(SearchRequest.builder().query("The World").topK(5).build());
assertThat(results).hasSize(3);
// Delete documents with country = BG
Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ,
new Filter.Key("country"), new Filter.Value("BG"));
store.delete(filterExpression);
results = store.similaritySearch(
SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
assertThat(results).hasSize(1);
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
}
});
}
@Test
void deleteWithStringFilterExpression() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createTestStore(context,
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),
new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) {
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
Map.of("country", "BG", "year", (short) 2020));
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", (short) 2023));
store.add(List.of(bgDocument, nlDocument, bgDocument2));
// Verify initial state
List<Document> results = store
.similaritySearch(SearchRequest.builder().query("The World").topK(5).build());
assertThat(results).hasSize(3);
store.delete("country == 'BG'");
results = store.similaritySearch(
SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
assertThat(results).hasSize(1);
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
}
});
}
@Test
void deleteWithComplexFilterExpression() {
this.contextRunner.run(context -> {
try (CassandraVectorStore store = createTestStore(context,
new SchemaColumn("type", DataTypes.TEXT, SchemaColumnTags.INDEXED),
new SchemaColumn("priority", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) {
var doc1 = new Document("Content 1", Map.of("type", "A", "priority", (short) 1));
var doc2 = new Document("Content 2", Map.of("type", "A", "priority", (short) 2));
var doc3 = new Document("Content 3", Map.of("type", "B", "priority", (short) 1));
store.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((short) 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);
store.delete(complexFilter);
var results = store.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 -> ((Short) doc.getMetadata().get("priority")).intValue())
.collect(Collectors.toList())).containsExactlyInAnyOrder(1, 1);
}
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {