[fix][dingo-executor] Fix disk ann index status issue (#1328)

dingodb · Jan 8, 2025 · a54de00 · a54de00
1 parent 1874910
commit a54de00
Showing 1 changed file with 22 additions and 7 deletions.
diff --git a/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/service/StoreService.java b/dingo-store-proxy/src/main/java/io/dingodb/store/proxy/service/StoreService.java
@@ -598,7 +598,7 @@ public List<VectorSearchResponse> vectorSearch(
                 .search(search)
                 .build();
             if (isDiskAnn) {
-                String diskAnnStatus = diskAnnStatus(requestTs, indexId);
+                String diskAnnStatus = diskAnnDetailStatus(requestTs, indexId);
                 if (diskAnnStatus != null && diskAnnStatus.equals("NODATA")) {
                     parameter.setUseBruteForce(true);
                 }
@@ -686,12 +686,16 @@ public String diskAnnBuild(long requestTs, CommonId indexId, long ts) {
             );
             Optional<VectorStateParameter> optionalState = Optional.ofNullable(vectorBuildResponse.getState());
             return optionalState
-                .map(state -> ((StateDiskAnnParam) state.getState()).getState().toString())
-                .orElse("INITIALIZED");
+                .map(state -> ((StateDiskAnnParam) state.getState()).getDiskannState().toString())
+                .orElse("DISKANN_INITIALIZED");
         }
 
         @Override
         public String diskAnnLoad(long requestTs, CommonId indexId, int nodesCacheNum, boolean warmup) {
+            String diskAnnStatus = diskAnnDetailStatus(requestTs, indexId);
+            if (diskAnnStatus != null && diskAnnStatus.equals("NODATA")) {
+                return "NODATA";
+            }
             VectorLoadResponse vectorLoadResponse = indexService(indexId, regionId).vectorLoad(
                 requestTs,
                 VectorLoadRequest.builder()
@@ -708,12 +712,23 @@ public String diskAnnLoad(long requestTs, CommonId indexId, int nodesCacheNum, b
             );
             Optional<VectorStateParameter> optionalState = Optional.ofNullable(vectorLoadResponse.getState());
             return optionalState
-                .map(state -> ((StateDiskAnnParam) state.getState()).getState().toString())
-                .orElse("INITIALIZED");
+                .map(state -> ((StateDiskAnnParam) state.getState()).getDiskannState().toString())
+                .orElse("DISKANN_INITIALIZED");
         }
 
         @Override
         public String diskAnnStatus(long requestTs, CommonId indexId) {
+            VectorStatusResponse vectorStatusResponse = indexService(indexId, regionId).vectorStatus(
+                requestTs,
+                VectorStatusRequest.builder().build()
+            );
+            Optional<VectorStateParameter> optionalState = Optional.ofNullable(vectorStatusResponse.getState());
+            return optionalState
+                .map(state -> ((StateDiskAnnParam) state.getState()).getDiskannState().toString())
+                .orElse("DISKANN_INITIALIZED");
+        }
+
+        public String diskAnnDetailStatus(long requestTs, CommonId indexId) {
             VectorStatusResponse vectorStatusResponse = indexService(indexId, regionId).vectorStatus(
                 requestTs,
                 VectorStatusRequest.builder().build()
@@ -732,8 +747,8 @@ public String diskAnnReset(long requestTs, CommonId indexId) {
             );
             Optional<VectorStateParameter> optionalState = Optional.ofNullable(vectorResetResponse.getState());
             return optionalState
-                .map(state -> ((StateDiskAnnParam) state.getState()).getState().toString())
-                .orElse("INITIALIZED");
+                .map(state -> ((StateDiskAnnParam) state.getState()).getDiskannState().toString())
+                .orElse("DISKANN_INITIALIZED");
         }
 
         @Override