From 6710318686ecd296bac10a5f991cca8a4ccafc7d Mon Sep 17 00:00:00 2001
From: presburger <yusheng.ma@zilliz.com>
Date: Mon, 29 Jul 2024 17:47:41 +0800
Subject: [PATCH] fix cagra cpu search bug (#736)

Signed-off-by: yusheng.ma <yusheng.ma@zilliz.com>
---
 src/index/gpu_raft/gpu_raft_cagra.cc | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/index/gpu_raft/gpu_raft_cagra.cc b/src/index/gpu_raft/gpu_raft_cagra.cc
index 49da15781..240052eef 100644
--- a/src/index/gpu_raft/gpu_raft_cagra.cc
+++ b/src/index/gpu_raft/gpu_raft_cagra.cc
@@ -38,7 +38,7 @@ class GpuRaftCagraHybridIndexNode : public GpuRaftCagraIndexNode<DataType> {
     Status
     Train(const DataSetPtr dataset, const Config& cfg) override {
         const GpuRaftCagraConfig& cagra_cfg = static_cast<const GpuRaftCagraConfig&>(cfg);
-        if (cagra_cfg.adapt_for_cpu)
+        if (cagra_cfg.adapt_for_cpu.value())
             adapt_for_cpu = true;
         return GpuRaftCagraIndexNode<DataType>::Train(dataset, cfg);
     }
@@ -111,6 +111,16 @@ class GpuRaftCagraHybridIndexNode : public GpuRaftCagraIndexNode<DataType> {
         return result;
     }
 
+    int64_t
+    Count() const override {
+        if (!adapt_for_cpu)
+            return GpuRaftCagraIndexNode<DataType>::Count();
+        if (!hnsw_index_) {
+            return 0;
+        }
+        return hnsw_index_->cur_element_count;
+    }
+
     Status
     Deserialize(const BinarySet& binset, const Config& config) override {
         if (binset.Contains(std::string(this->Type()) + "_cpu")) {