Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

VectoriaDB server fixes #127

Merged
merged 3 commits into from
Feb 7, 2024
Merged

VectoriaDB server fixes #127

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d785d2a
remove the index folder in the SearchMode.dropIndex(...)
kirillvasilenko Feb 6, 2024
00e439c
fix handling the case when the context is cancelled
kirillvasilenko Feb 6, 2024
e666127
complete requestObserver even if there is no progressIndicator
kirillvasilenko Feb 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions vectoriadb-java-client/src/main/java/jetbrains/vectoriadb/client/VectoriaDBClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,10 @@ public void onCompleted() {
progressIndicator.accept(-1, -1);
requestObserver.onCompleted();
progressIndicator.accept(Integer.MAX_VALUE, Integer.MAX_VALUE);
} else {
requestObserver.onCompleted();
}

try {
finishedLatch.await();
} catch (InterruptedException e) {
Expand Down
10 changes: 7 additions & 3 deletions vectoriadb-server/src/main/java/jetbrains/vectoriadb/server/IndexManagerServiceImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -902,6 +902,8 @@ public void dropIndex(IndexManagerOuterClass.IndexNameRequest request, StreamObs
@SuppressWarnings("resource") var indexReader = fetchIndexReader(indexName);
indexReader.deleteIndex();

var indexDir = indexMetadatas.get(indexName).dir;
FileUtils.deleteDirectory(indexDir.toFile());
//noinspection resource
indexReaders.remove(indexName);
indexStates.remove(indexName);
Expand All @@ -917,7 +919,7 @@ public void dropIndex(IndexManagerOuterClass.IndexNameRequest request, StreamObs

@NotNull
private IndexReader fetchIndexReader(final String indexName) {
return indexReaders.computeIfAbsent(indexName, r -> {
return indexReaders.computeIfAbsent(indexName, _ -> {
var metadata = indexMetadatas.get(indexName);
return new IndexReader(indexName, dimensions, maxConnectionsPerVertex, maxCandidatesReturned,
compressionRatio, metadata.dir, metadata.distance, diskCache);
Expand Down Expand Up @@ -1029,7 +1031,7 @@ public void createIndex(IndexManagerOuterClass.CreateIndexRequest request,
public void buildIndex(IndexManagerOuterClass.IndexNameRequest request, StreamObserver<Empty> responseObserver) {
try {
var indexName = request.getIndexName();
var indexState = indexStates.compute(indexName, (k, state) -> {
var indexState = indexStates.compute(indexName, (_, state) -> {
if (state == IndexState.UPLOADED || state == IndexState.CREATED) {
return IndexState.IN_BUILD_QUEUE;
} else {
Expand Down Expand Up @@ -1295,7 +1297,7 @@ public void dropIndex(IndexManagerOuterClass.IndexNameRequest request, StreamObs
try {
var indexName = request.getIndexName();

indexStates.compute(indexName, (k, state) -> {
indexStates.compute(indexName, (_, state) -> {
if (state == IndexState.CREATED || state == IndexState.BUILT || state == IndexState.UPLOADED) {
return IndexState.BROKEN;
} else {
Expand Down Expand Up @@ -1365,11 +1367,13 @@ public ServiceIndexBuildProgressListener(StreamObserver<IndexManagerOuterClass.B
public void progress(IndexBuildProgressInfo progressInfo) {
if (context.isCancelled()) {
try {
progressTracker.removeListener(this);
responseObserver.onCompleted();
} catch (Exception e) {
progressTracker.removeListener(this);
responseObserver.onError(new StatusRuntimeException(Status.INTERNAL.withCause(e)));
}
return;
}

try {
Expand Down
41 changes: 34 additions & 7 deletions vectoriadb-server/src/test/java/jetbrains/vectoriadb/server/IndexManagerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,7 @@ public void testListIndexes() throws Exception {

indexes = listIndexes(indexManagerService);
Assert.assertEquals(1, indexes.size());
Assert.assertEquals(indexName + 1, indexes.get(0));
Assert.assertEquals(indexName + 1, indexes.getFirst());

createIndex(indexName + 2, indexManagerService, IndexManagerOuterClass.Distance.L2);
indexes = listIndexes(indexManagerService);
Expand Down Expand Up @@ -247,6 +247,26 @@ public void testDropSearchIndex() throws Exception {
});
}

@Test
public void testDropIndexInSearchModeAndCreateItAgain() throws Exception {
var indexOne = "testDropSearchIndexOne";

executeInServiceContext(indexManagerService -> {
generateIndex(indexOne, L2DistanceFunction.INSTANCE, 64, 1, indexManagerService);

switchToSearchMode(indexManagerService);

dropIndex(indexOne, indexManagerService);

var indexes = listIndexes(indexManagerService);
Assert.assertEquals(0, indexes.size());

switchToBuildMode(indexManagerService);

generateIndex(indexOne, L2DistanceFunction.INSTANCE, 64, 1, indexManagerService);
});
}

@Test
public void testTwoIndexesSimultaniously() throws Exception {
var buildDir = System.getProperty("exodus.tests.buildDirectory");
Expand Down Expand Up @@ -471,7 +491,7 @@ private static List<String> listIndexes(IndexManagerServiceImpl indexManagerServ
indexManagerService.listIndexes(Empty.newBuilder().build(), listIndexesRecorder);

checkCompleteness(listIndexesRecorder);
return listIndexesRecorder.getValues().get(0).getIndexNamesList();
return listIndexesRecorder.getValues().getFirst().getIndexNamesList();
}


Expand All @@ -492,7 +512,7 @@ private static void buildIndex(String indexName, IndexManagerServiceImpl indexMa

checkCompleteness(indexStateRecorder);

var response = indexStateRecorder.getValues().get(0);
var response = indexStateRecorder.getValues().getFirst();
var indexState = response.getState();
if (indexState == IndexManagerOuterClass.IndexState.BUILDING ||
indexState == IndexManagerOuterClass.IndexState.IN_BUILD_QUEUE) {
Expand Down Expand Up @@ -531,6 +551,13 @@ private static void switchToSearchMode(IndexManagerServiceImpl indexManagerServi
checkCompleteness(switchToSearchModeRecorder);
}

private static void switchToBuildMode(IndexManagerServiceImpl indexManagerService) throws Exception {
var streamRecorder = StreamRecorder.<Empty>create();
indexManagerService.switchToBuildMode(Empty.newBuilder().build(), streamRecorder);

checkCompleteness(streamRecorder);
}

private static byte[][] findNearestNeighbours(IndexManagerServiceImpl indexManagerService,
float[] queryVector, String indexName, int k) throws Exception {
var findNearestVectorsRecorder = StreamRecorder.<IndexManagerOuterClass.FindNearestNeighboursResponse>create();
Expand All @@ -546,7 +573,7 @@ private static byte[][] findNearestNeighbours(IndexManagerServiceImpl indexManag

checkCompleteness(findNearestVectorsRecorder);

var response = findNearestVectorsRecorder.getValues().get(0);
var response = findNearestVectorsRecorder.getValues().getFirst();
var nearestVectors = response.getIdsList();
var result = new byte[nearestVectors.size()][];

Expand Down Expand Up @@ -674,9 +701,9 @@ private static void searchNeighbours(String indexName, int vectorsCount, int vec
@NotNull
private static IndexManagerOuterClass.Distance convertDistanceFunction(DistanceFunction distanceFunction) {
return switch (distanceFunction) {
case L2DistanceFunction l2DistanceFunction -> IndexManagerOuterClass.Distance.L2;
case DotDistanceFunction dotDistanceFunction -> IndexManagerOuterClass.Distance.DOT;
case CosineDistanceFunction cosineDistanceFunction -> IndexManagerOuterClass.Distance.COSINE;
case L2DistanceFunction _ -> IndexManagerOuterClass.Distance.L2;
case DotDistanceFunction _ -> IndexManagerOuterClass.Distance.DOT;
case CosineDistanceFunction _ -> IndexManagerOuterClass.Distance.COSINE;
case null, default -> throw new IllegalArgumentException("Unknown distance function " + distanceFunction);
};
}
Expand Down
Loading