fix cagra cpu search bug (#736)

Signed-off-by: yusheng.ma <[email protected]>
zilliztech · Jul 29, 2024 · 6710318 · 6710318
1 parent 752e5e0
commit 6710318
Showing 1 changed file with 11 additions and 1 deletion.
diff --git 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")) {