Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix PgVectorStore doDelete function as batch

Browse Source 
Signed-off-by: CChuYong <yeongmin1061@gmail.com>
This commit is contained in:
CChuYong
2025-03-12 16:44:22 +09:00
committed by Ilayaperumal Gopinathan
parent d5203ed038
commit 268248ba3c
2 changed files with 60 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java
View File

@@ -152,6 +152,7 @@ import org.springframework.util.StringUtils;
* @author Soby Chacko
* @author Sebastien Deleuze
* @author Jihoon Kim
* @author YeongMin Song
* @since 1.0.0
*/
public class PgVectorStore extends AbstractObservationVectorStore implements InitializingBean {
@@ -319,12 +320,21 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
@Override
public void doDelete(List<String> idList) {
int updateCount = 0;
for (String id : idList) {
int count = this.jdbcTemplate.update("DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?",
UUID.fromString(id));
updateCount = updateCount + count;
}
String sql = "DELETE FROM " + getFullyQualifiedTableName() + " WHERE id = ?";
this.jdbcTemplate.batchUpdate(sql, new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement ps, int i) throws SQLException {
var id = idList.get(i);
StatementCreatorUtils.setParameterValue(ps, 1, SqlTypeValue.TYPE_UNKNOWN, convertIdToPgType(id));
}
@Override
public int getBatchSize() {
return idList.size();
}
});
}
@Override

44
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
View File

@@ -75,6 +75,7 @@ import static org.assertj.core.api.Assertions.assertThat;
* @author Christian Tzolov
* @author Thomas Vitale
* @author Jihoon Kim
* @author YeongMin Song
*/
@Testcontainers
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -232,6 +233,47 @@ public class PgVectorStoreIT extends BaseVectorStoreTests {
});
}
@Test
public void testBulkOperationWithUuidIdType() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + "COSINE_DISTANCE")
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
List<Document> documents = List.of(
new Document(new RandomIdGenerator().generateId(), "TEXT", new HashMap<>()),
new Document(new RandomIdGenerator().generateId(), "TEXT", new HashMap<>()),
new Document(new RandomIdGenerator().generateId(), "TEXT", new HashMap<>()));
vectorStore.add(documents);
List<String> idList = documents.stream().map(Document::getId).toList();
vectorStore.delete(idList);
dropTable(context);
});
}
@Test
public void testBulkOperationWithNonUuidIdType() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.pgvector.distanceType=" + "COSINE_DISTANCE")
.withPropertyValues("test.spring.ai.vectorstore.pgvector.initializeSchema=" + false)
.withPropertyValues("test.spring.ai.vectorstore.pgvector.idType=" + "TEXT")
.run(context -> {
VectorStore vectorStore = context.getBean(VectorStore.class);
initSchema(context);
List<Document> documents = List.of(new Document("NON_UUID_1", "TEXT", new HashMap<>()),
new Document("NON_UUID_2", "TEXT", new HashMap<>()),
new Document("NON_UUID_3", "TEXT", new HashMap<>()));
vectorStore.add(documents);
List<String> idList = documents.stream().map(Document::getId).toList();
vectorStore.delete(idList);
dropTable(context);
});
}
@ParameterizedTest(name = "Filter expression {0} should return {1} records ")
@MethodSource("provideFilters")
public void searchWithInFilter(String expression, Integer expectedRecords) {
@@ -436,6 +478,8 @@ public class PgVectorStoreIT extends BaseVectorStoreTests {
PgVectorStore vectorStore = context.getBean(PgVectorStore.class);
Optional<JdbcTemplate> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
dropTable(context);
});
}