Skip to content

Commit 71e0e36

Browse files
committed
Add missing integration tests for delete by ID API in vector store implementations.
Signed-off-by: Soby Chacko <[email protected]>
1 parent 015662a commit 71e0e36

File tree

12 files changed

+379
-0
lines changed

12 files changed

+379
-0
lines changed

vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,38 @@ void searchWithThreshold() {
421421
});
422422
}
423423

424+
@Test
425+
void deleteById() {
426+
this.contextRunner.run(context -> {
427+
try (CassandraVectorStore store = createTestStore(context,
428+
new SchemaColumn("country", DataTypes.TEXT, SchemaColumnTags.INDEXED),
429+
new SchemaColumn("year", DataTypes.SMALLINT, SchemaColumnTags.INDEXED))) {
430+
431+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
432+
Map.of("country", "BG", "year", (short) 2020));
433+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
434+
Map.of("country", "NL"));
435+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
436+
Map.of("country", "BG", "year", (short) 2023));
437+
438+
store.add(List.of(bgDocument, nlDocument, bgDocument2));
439+
440+
// Verify initial state
441+
List<Document> results = store
442+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).build());
443+
assertThat(results).hasSize(3);
444+
445+
store.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
446+
447+
results = store.similaritySearch(
448+
SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
449+
450+
assertThat(results).hasSize(1);
451+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
452+
}
453+
});
454+
}
455+
424456
@Test
425457
void deleteByFilter() {
426458
this.contextRunner.run(context -> {

vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,41 @@ public void addAndSearchWithFilters() {
169169
});
170170
}
171171

172+
@Test
173+
public void deleteById() {
174+
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> {
175+
VectorStore vectorStore = context.getBean(VectorStore.class);
176+
177+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
178+
Map.of("country", "BG", "year", 2020));
179+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
180+
Map.of("country", "NL", "year", 2021));
181+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
182+
Map.of("country", "BG", "year", 2023));
183+
184+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
185+
186+
SearchRequest searchRequest = SearchRequest.builder()
187+
.query("The World")
188+
.topK(5)
189+
.similarityThresholdAll()
190+
.build();
191+
192+
List<Document> results = vectorStore.similaritySearch(searchRequest);
193+
assertThat(results).hasSize(3);
194+
195+
Filter.Expression filterExpression = new Filter.Expression(Filter.ExpressionType.EQ,
196+
new Filter.Key("country"), new Filter.Value("BG"));
197+
198+
vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
199+
200+
// Verify deletion - should only have NL document remaining
201+
results = vectorStore.similaritySearch(searchRequest);
202+
assertThat(results).hasSize(1);
203+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
204+
});
205+
}
206+
172207
@Test
173208
public void deleteWithFilterExpression() {
174209
this.contextRunner.run(context -> {

vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,40 @@ public void searchWithThreshold(String distanceType) {
362362
});
363363
}
364364

365+
@Test
366+
public void deleteById() {
367+
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> {
368+
VectorStore vectorStore = context.getBean(VectorStore.class);
369+
370+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
371+
Map.of("country", "BG", "year", 2020));
372+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
373+
Map.of("country", "NL", "year", 2021));
374+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
375+
Map.of("country", "BG", "year", 2023));
376+
377+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
378+
379+
SearchRequest searchRequest = SearchRequest.builder()
380+
.query("The World")
381+
.topK(5)
382+
.similarityThresholdAll()
383+
.build();
384+
385+
List<Document> results = vectorStore.similaritySearch(searchRequest);
386+
assertThat(results).hasSize(3);
387+
388+
vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
389+
390+
// Verify deletion - should only have NL document remaining
391+
results = vectorStore.similaritySearch(searchRequest);
392+
assertThat(results).hasSize(1);
393+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
394+
395+
dropTable(context);
396+
});
397+
}
398+
365399
@Test
366400
public void deleteByFilter() {
367401
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.mariadb.distanceType=COSINE").run(context -> {

vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,40 @@ public void searchWithThreshold(String metricType) {
278278
});
279279
}
280280

281+
@Test
282+
public void deleteById() {
283+
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=COSINE").run(context -> {
284+
VectorStore vectorStore = context.getBean(VectorStore.class);
285+
286+
resetCollection(vectorStore);
287+
288+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
289+
Map.of("country", "BG", "year", 2020));
290+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
291+
Map.of("country", "NL", "year", 2021));
292+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
293+
Map.of("country", "BG", "year", 2023));
294+
295+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
296+
297+
SearchRequest searchRequest = SearchRequest.builder()
298+
.query("The World")
299+
.topK(5)
300+
.similarityThresholdAll()
301+
.build();
302+
303+
List<Document> results = vectorStore.similaritySearch(searchRequest);
304+
assertThat(results).hasSize(3);
305+
306+
vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
307+
308+
// Verify deletion - should only have NL document remaining
309+
results = vectorStore.similaritySearch(searchRequest);
310+
assertThat(results).hasSize(1);
311+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
312+
});
313+
}
314+
281315
@Test
282316
public void deleteByFilter() {
283317
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=COSINE").run(context -> {

vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,39 @@ public void searchWithThreshold() {
257257
});
258258
}
259259

260+
@Test
261+
void deleteById() {
262+
this.contextRunner.run(context -> {
263+
VectorStore vectorStore = context.getBean(VectorStore.class);
264+
265+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
266+
Map.of("country", "BG", "year", 2020));
267+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
268+
Map.of("country", "NL", "year", 2021));
269+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
270+
Map.of("country", "BG", "year", 2023));
271+
272+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
273+
Thread.sleep(5000); // Wait for indexing
274+
275+
SearchRequest searchRequest = SearchRequest.builder()
276+
.query("The World")
277+
.topK(5)
278+
.similarityThresholdAll()
279+
.build();
280+
281+
List<Document> results = vectorStore.similaritySearch(searchRequest);
282+
assertThat(results).hasSize(3);
283+
284+
vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
285+
Thread.sleep(1000); // Wait for deletion to be processed
286+
287+
results = vectorStore.similaritySearch(searchRequest);
288+
assertThat(results).hasSize(1);
289+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
290+
});
291+
}
292+
260293
@Test
261294
void deleteByFilter() {
262295
this.contextRunner.run(context -> {

vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,37 @@ void ensureIdIndexGetsCreated() {
305305
.isTrue());
306306
}
307307

308+
@Test
309+
void deleteById() {
310+
this.contextRunner.run(context -> {
311+
VectorStore vectorStore = context.getBean(VectorStore.class);
312+
313+
var bgDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
314+
Map.of("country", "BG", "year", 2020));
315+
var nlDocument = new Document("The World is Big and Salvation Lurks Around the Corner",
316+
Map.of("country", "NL", "year", 2021));
317+
var bgDocument2 = new Document("The World is Big and Salvation Lurks Around the Corner",
318+
Map.of("country", "BG", "year", 2023));
319+
320+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
321+
322+
SearchRequest searchRequest = SearchRequest.builder()
323+
.query("The World")
324+
.topK(5)
325+
.similarityThresholdAll()
326+
.build();
327+
328+
List<Document> results = vectorStore.similaritySearch(searchRequest);
329+
assertThat(results).hasSize(3);
330+
331+
vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
332+
333+
results = vectorStore.similaritySearch(searchRequest);
334+
assertThat(results).hasSize(1);
335+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
336+
});
337+
}
338+
308339
@Test
309340
void deleteByFilter() {
310341
this.contextRunner.run(context -> {

vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,38 @@ public void searchDocumentsInTwoIndicesTest() {
418418
});
419419
}
420420

421+
@Test
422+
void deleteById() {
423+
getContextRunner().run(context -> {
424+
OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class);
425+
426+
var bgDocument = new Document("1", "The World is Big and Salvation Lurks Around the Corner",
427+
Map.of("country", "BG", "year", 2020, "activationDate", new Date(1000)));
428+
var nlDocument = new Document("2", "The World is Big and Salvation Lurks Around the Corner",
429+
Map.of("country", "NL", "activationDate", new Date(2000)));
430+
var bgDocument2 = new Document("3", "The World is Big and Salvation Lurks Around the Corner",
431+
Map.of("country", "BG", "year", 2023, "activationDate", new Date(3000)));
432+
433+
vectorStore.add(List.of(bgDocument, nlDocument, bgDocument2));
434+
435+
Awaitility.await()
436+
.until(() -> vectorStore.similaritySearch(SearchRequest.builder().query("The World").topK(5).build()),
437+
hasSize(3));
438+
439+
vectorStore.delete(List.of(bgDocument.getId(), bgDocument2.getId()));
440+
441+
Awaitility.await()
442+
.until(() -> vectorStore.similaritySearch(SearchRequest.builder().query("The World").topK(5).build()),
443+
hasSize(1));
444+
445+
List<Document> results = vectorStore
446+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
447+
448+
assertThat(results).hasSize(1);
449+
assertThat(results.get(0).getMetadata()).containsEntry("country", "NL");
450+
});
451+
}
452+
421453
@Test
422454
void deleteByFilter() {
423455
getContextRunner().run(context -> {

vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -317,6 +317,42 @@ public void searchWithThreshold(String distanceType) {
317317
});
318318
}
319319

320+
@Test
321+
void deleteById() {
322+
this.contextRunner
323+
.withPropertyValues("test.spring.ai.vectorstore.oracle.distanceType=COSINE",
324+
"test.spring.ai.vectorstore.oracle.searchAccuracy=" + OracleVectorStore.DEFAULT_SEARCH_ACCURACY)
325+
.run(context -> {
326+
VectorStore vectorStore = context.getBean(VectorStore.class);
327+
328+
var doc1 = new Document("The World is Big and Salvation Lurks Around the Corner",
329+
Map.of("country", "BG", "year", 2020));
330+
var doc2 = new Document("The World is Big and Salvation Lurks Around the Corner",
331+
Map.of("country", "NL"));
332+
var doc3 = new Document("The World is Big and Salvation Lurks Around the Corner",
333+
Map.of("country", "BG", "year", 2023));
334+
335+
vectorStore.add(List.of(doc1, doc2, doc3));
336+
337+
// Delete first two documents
338+
vectorStore.delete(List.of(doc1.getId(), doc2.getId()));
339+
340+
List<Document> results = vectorStore.similaritySearch(
341+
SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
342+
343+
assertThat(results).hasSize(1);
344+
assertThat(results.get(0).getId()).isEqualTo(doc3.getId());
345+
assertThat(results).hasSize(1);
346+
assertThat(results.get(0).getId()).isEqualTo(doc3.getId());
347+
assertThat(results.get(0).getMetadata())
348+
.hasEntrySatisfying("country", value -> assertThat(value.toString()).isEqualTo("\"BG\""))
349+
.hasEntrySatisfying("year", value -> assertThat(value.toString()).isEqualTo("2023"));
350+
351+
// Clean up remaining document
352+
vectorStore.delete(List.of(doc3.getId()));
353+
});
354+
}
355+
320356
@Test
321357
void deleteByFilter() {
322358
this.contextRunner

vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -260,6 +260,34 @@ public void searchThresholdTest() {
260260
});
261261
}
262262

263+
@Test
264+
void deleteById() {
265+
this.contextRunner.run(context -> {
266+
VectorStore vectorStore = context.getBean(VectorStore.class);
267+
268+
var doc1 = new Document("The World is Big and Salvation Lurks Around the Corner",
269+
Map.of("country", "BG", "year", 2020));
270+
var doc2 = new Document("The World is Big and Salvation Lurks Around the Corner", Map.of("country", "NL"));
271+
var doc3 = new Document("The World is Big and Salvation Lurks Around the Corner",
272+
Map.of("country", "BG", "year", 2023));
273+
274+
vectorStore.add(List.of(doc1, doc2, doc3));
275+
276+
// Delete first two documents
277+
vectorStore.delete(List.of(doc1.getId(), doc2.getId()));
278+
279+
List<Document> results = vectorStore
280+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
281+
282+
assertThat(results).hasSize(1);
283+
assertThat(results.get(0).getId()).isEqualTo(doc3.getId());
284+
assertThat(results.get(0).getMetadata()).containsEntry("country", "BG").containsEntry("year", 2023L);
285+
286+
// Clean up remaining document
287+
vectorStore.delete(List.of(doc3.getId()));
288+
});
289+
}
290+
263291
@Test
264292
void deleteByFilter() {
265293
this.contextRunner.run(context -> {

vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,34 @@ void searchWithThreshold() {
264264
});
265265
}
266266

267+
@Test
268+
void deleteById() {
269+
this.contextRunner.run(context -> {
270+
VectorStore vectorStore = context.getBean(VectorStore.class);
271+
272+
var doc1 = new Document("The World is Big and Salvation Lurks Around the Corner",
273+
Map.of("country", "BG", "year", 2020));
274+
var doc2 = new Document("The World is Big and Salvation Lurks Around the Corner", Map.of("country", "NL"));
275+
var doc3 = new Document("The World is Big and Salvation Lurks Around the Corner",
276+
Map.of("country", "BG", "year", 2023));
277+
278+
vectorStore.add(List.of(doc1, doc2, doc3));
279+
280+
// Delete first two documents
281+
vectorStore.delete(List.of(doc1.getId(), doc2.getId()));
282+
283+
List<Document> results = vectorStore
284+
.similaritySearch(SearchRequest.builder().query("The World").topK(5).similarityThresholdAll().build());
285+
286+
assertThat(results).hasSize(1);
287+
assertThat(results.get(0).getId()).isEqualTo(doc3.getId());
288+
assertThat(results.get(0).getMetadata()).containsEntry("country", "BG").containsEntry("year", "2023");
289+
290+
// Clean up remaining document
291+
vectorStore.delete(List.of(doc3.getId()));
292+
});
293+
}
294+
267295
@Test
268296
void deleteByFilter() {
269297
this.contextRunner.run(context -> {

0 commit comments

Comments
 (0)