From a54de00e509ef80c844dbbb98b4aa3b1b47feb63 Mon Sep 17 00:00:00 2001
From: githubgxll <1094462054@qq.com>
Date: Wed, 8 Jan 2025 10:01:24 +0800
Subject: [PATCH] [fix][dingo-executor] Fix disk ann index status issue (#1328)

---
 .../store/proxy/service/StoreService.java     | 29 ++++++++++++++-----
 1 file changed, 22 insertions(+), 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
index 10ee6694c4..9ecda5ed9b 100644
--- 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