fix MapOffsetToUid (#4493)

Signed-off-by: Nshengjun.li <shengjun.li@zilliz.com>

core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIDMAP.cpp

@@ -51,6 +51,8 @@ BinaryIDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config) {
     auto p_dist = (float*)malloc(p_dist_size);
 
     QueryImpl(rows, (uint8_t*)p_data, k, p_dist, p_id, config);
+    MapOffsetToUid(p_id, static_cast<size_t>(elems));
+
     auto ret_ds = std::make_shared<Dataset>();
     ret_ds->Set(meta::IDS, p_id);
     ret_ds->Set(meta::DISTANCE, p_dist);
@@ -122,8 +124,6 @@ BinaryIDMAP::QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distanc
     }
 
     index_->metric_type = default_type;
-
-    MapOffsetToUid(labels, static_cast<size_t>(n * k));
 }
 
 }  // namespace knowhere

core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp

@@ -58,6 +58,7 @@ BinaryIVF::Query(const DatasetPtr& dataset_ptr, const Config& config) {
         auto p_dist = (float*)malloc(p_dist_size);
 
         QueryImpl(rows, (uint8_t*)p_data, k, p_dist, p_id, config);
+        MapOffsetToUid(p_id, static_cast<size_t>(elems));
 
         auto ret_ds = std::make_shared<Dataset>();
         ret_ds->Set(meta::IDS, p_id);
@@ -187,8 +188,6 @@ BinaryIVF::QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances
             distances[i] = static_cast<float>(i_distances[i]);
         }
     }
-
-    MapOffsetToUid(labels, static_cast<size_t>(n * k));
 }
 
 }  // namespace knowhere

core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp

@@ -82,6 +82,8 @@ IDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config) {
     auto p_dist = (float*)malloc(p_dist_size);
 
     QueryImpl(rows, (float*)p_data, k, p_dist, p_id, config);
+    MapOffsetToUid(p_id, static_cast<size_t>(elems));
+
     auto ret_ds = std::make_shared<Dataset>();
     ret_ds->Set(meta::IDS, p_id);
     ret_ds->Set(meta::DISTANCE, p_dist);
@@ -139,8 +141,6 @@ IDMAP::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int6
         index_->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
     index_->search(n, (float*)data, k, distances, labels, GetBlacklist());
     index_->metric_type = default_type;
-
-    MapOffsetToUid(labels, static_cast<size_t>(n * k));
 }
 
 }  // namespace knowhere

core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp

@@ -104,6 +104,7 @@ IVF::Query(const DatasetPtr& dataset_ptr, const Config& config) {
         auto p_dist = (float*)malloc(p_dist_size);
 
         QueryImpl(rows, (float*)p_data, k, p_dist, p_id, config);
+        MapOffsetToUid(p_id, static_cast<size_t>(elems));
 
         auto ret_ds = std::make_shared<Dataset>();
         ret_ds->Set(meta::IDS, p_id);
@@ -311,8 +312,6 @@ IVF::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_
                         << ", data search cost: " << faiss::indexIVF_stats.search_time;
     faiss::indexIVF_stats.quantization_time = 0;
     faiss::indexIVF_stats.search_time = 0;
-
-    MapOffsetToUid(labels, static_cast<size_t>(n * k));
 }
 
 void

core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIDMAP.cpp

@@ -105,8 +105,6 @@ GPUIDMAP::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, i
         index_->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
     index_->search(n, (float*)data, k, distances, labels, GetBlacklist());
     index_->metric_type = default_type;
-
-    MapOffsetToUid(labels, static_cast<size_t>(n * k));
 }
 
 void

core/src/index/knowhere/knowhere/index/vector_index/gpu/IndexGPUIVF.cpp

@@ -148,8 +148,6 @@ GPUIVF::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int
             device_index->search(search_size, (float*)data + i * dim, k, distances + i * k, labels + i * k,
                                  GetBlacklist());
         }
-
-        MapOffsetToUid(labels, static_cast<size_t>(n * k));
     } else {
         KNOWHERE_THROW_MSG("Not a GpuIndexIVF type.");
     }

core/src/index/unittest/test_annoy.cpp

@@ -60,7 +60,7 @@ TEST_P(AnnoyTest, annoy_basic) {
         ASSERT_ANY_THROW(index_->Dim());
     }
 
-    index_->BuildAll(base_dataset, conf);  // Train + Add
+    index_->BuildAll(base_dataset, conf);  // Train + AddWithoutIds
     ASSERT_EQ(index_->Count(), nb);
     ASSERT_EQ(index_->Dim(), dim);
 
@@ -94,7 +94,7 @@ TEST_P(AnnoyTest, annoy_basic) {
 TEST_P(AnnoyTest, annoy_delete) {
     assert(!xb.empty());
 
-    index_->BuildAll(base_dataset, conf);  // Train + Add
+    index_->BuildAll(base_dataset, conf);  // Train + AddWithoutIds
     ASSERT_EQ(index_->Count(), nb);
     ASSERT_EQ(index_->Dim(), dim);