Fix GetVectorById Dim, Return Dataset and Set All HasRawData to False…

… Except HNSW (#957) Signed-off-by: Patrick Weizhi Xu <[email protected]>
milvus-io · Jun 26, 2023 · b6f01a6 · b6f01a6
1 parent b9bdbfc
commit b6f01a6
Show file tree

Hide file tree

Showing 20 changed files with 29 additions and 21 deletions.
diff --git a/knowhere/index/vector_index/IndexAnnoy.cpp b/knowhere/index/vector_index/IndexAnnoy.cpp
@@ -127,6 +127,7 @@ IndexAnnoy::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
     }
 
     GET_DATA_WITH_IDS(dataset_ptr)
+    auto dim = Dim();
 
     float* p_x = nullptr;
     try {
@@ -142,7 +143,7 @@ IndexAnnoy::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
         }
         KNOWHERE_THROW_MSG(e.what());
     }
-    return GenResultDataset(p_x);
+    return GenResultDataset(rows, dim, p_x);
 }
 
 DatasetPtr

diff --git a/knowhere/index/vector_index/IndexAnnoy.h b/knowhere/index/vector_index/IndexAnnoy.h
@@ -48,7 +48,7 @@ class IndexAnnoy : public VecIndex {
 
     bool
     HasRawData(const std::string& /*metric_type*/) const override {
-        return true;
+        return false;
     }
 
     DatasetPtr

diff --git a/knowhere/index/vector_index/IndexBinaryIDMAP.cpp b/knowhere/index/vector_index/IndexBinaryIDMAP.cpp
@@ -43,6 +43,7 @@ BinaryIDMAP::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config)
     }
 
     GET_DATA_WITH_IDS(dataset_ptr)
+    auto dim = Dim();
 
     uint8_t* p_x = nullptr;
     auto release_when_exception = [&]() {
@@ -59,7 +60,7 @@ BinaryIDMAP::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config)
             KNOWHERE_THROW_IF_NOT_FMT(id >= 0 && id < bin_idmap_index->ntotal, "invalid id %ld", id);
             bin_idmap_index->reconstruct(id, p_x + i * dim / 8);
         }
-        return GenResultDataset(p_x);
+        return GenResultDataset(rows, dim, p_x);
     } catch (faiss::FaissException& e) {
         release_when_exception();
         KNOWHERE_THROW_MSG(e.what());

diff --git a/knowhere/index/vector_index/IndexBinaryIDMAP.h b/knowhere/index/vector_index/IndexBinaryIDMAP.h
@@ -48,7 +48,7 @@ class BinaryIDMAP : public VecIndex, public FaissBaseBinaryIndex {
 
     bool
     HasRawData(const std::string& /*metric_type*/) const override {
-        return true;
+        return false;
     }
 
     DatasetPtr

diff --git a/knowhere/index/vector_index/IndexBinaryIVF.cpp b/knowhere/index/vector_index/IndexBinaryIVF.cpp
@@ -54,6 +54,7 @@ BinaryIVF::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
     }
 
     GET_DATA_WITH_IDS(dataset_ptr)
+    auto dim = Dim();
 
     uint8_t* p_x = nullptr;
     auto release_when_exception = [&]() {
@@ -71,7 +72,7 @@ BinaryIVF::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
             KNOWHERE_THROW_IF_NOT_FMT(id >= 0 && id < bin_ivf_index->ntotal, "invalid id %ld", id);
             bin_ivf_index->reconstruct(id, p_x + i * dim / 8);
         }
-        return GenResultDataset(p_x);
+        return GenResultDataset(rows, dim, p_x);
     } catch (faiss::FaissException& e) {
         release_when_exception();
         KNOWHERE_THROW_MSG(e.what());

diff --git a/knowhere/index/vector_index/IndexBinaryIVF.h b/knowhere/index/vector_index/IndexBinaryIVF.h
@@ -53,7 +53,7 @@ class BinaryIVF : public VecIndex, public FaissBaseBinaryIndex {
 
     bool
     HasRawData(const std::string& /*metric_type*/) const override {
-        return true;
+        return false;
     }
 
     DatasetPtr

diff --git a/knowhere/index/vector_index/IndexHNSW.cpp b/knowhere/index/vector_index/IndexHNSW.cpp
@@ -123,6 +123,7 @@ IndexHNSW::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
     }
 
     GET_DATA_WITH_IDS(dataset_ptr)
+    auto dim = Dim();
 
     float* p_x = nullptr;
     try {
@@ -138,7 +139,7 @@ IndexHNSW::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
         }
         KNOWHERE_THROW_MSG(e.what());
     }
-    return GenResultDataset(p_x);
+    return GenResultDataset(rows, dim, p_x);
 }
 
 DatasetPtr

diff --git a/knowhere/index/vector_index/IndexIDMAP.cpp b/knowhere/index/vector_index/IndexIDMAP.cpp
@@ -76,6 +76,7 @@ IDMAP::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
     }
 
     GET_DATA_WITH_IDS(dataset_ptr)
+    auto dim = Dim();
 
     float* p_x = nullptr;
     auto release_when_exception = [&]() {
@@ -92,7 +93,7 @@ IDMAP::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
             KNOWHERE_THROW_IF_NOT_FMT(id >= 0 && id < idmap_index->ntotal, "invalid id %ld", id);
             idmap_index->reconstruct(id, p_x + i * dim);
         }
-        return GenResultDataset(p_x);
+        return GenResultDataset(rows, dim, p_x);
     } catch (faiss::FaissException& e) {
         release_when_exception();
         KNOWHERE_THROW_MSG(e.what());

diff --git a/knowhere/index/vector_index/IndexIDMAP.h b/knowhere/index/vector_index/IndexIDMAP.h
@@ -48,7 +48,7 @@ class IDMAP : public VecIndex, public FaissBaseIndex {
 
     bool
     HasRawData(const std::string& /*metric_type*/) const override {
-        return true;
+        return false;
     }
 
     DatasetPtr

diff --git a/knowhere/index/vector_index/IndexIVF.cpp b/knowhere/index/vector_index/IndexIVF.cpp
@@ -98,6 +98,7 @@ IVF::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
     }
 
     GET_DATA_WITH_IDS(dataset_ptr)
+    auto dim = Dim();
 
     float* p_x = nullptr;
     auto release_when_exception = [&]() {
@@ -115,7 +116,7 @@ IVF::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
             KNOWHERE_THROW_IF_NOT_FMT(id >= 0 && id < ivf_index->ntotal, "invalid id %ld", id);
             ivf_index->reconstruct(id, p_x + i * dim);
         }
-        return GenResultDataset(p_x);
+        return GenResultDataset(rows, dim, p_x);
     } catch (faiss::FaissException& e) {
         release_when_exception();
         KNOWHERE_THROW_MSG(e.what());

diff --git a/knowhere/index/vector_index/IndexIVF.h b/knowhere/index/vector_index/IndexIVF.h
@@ -53,7 +53,7 @@ class IVF : public VecIndex, public FaissBaseIndex {
 
     bool
     HasRawData(const std::string& /*metric_type*/) const override {
-        return true;
+        return false;
     }
 
     DatasetPtr

diff --git a/knowhere/index/vector_index/adapter/VectorAdapter.cpp b/knowhere/index/vector_index/adapter/VectorAdapter.cpp
@@ -36,8 +36,10 @@ GenDatasetWithIds(const int64_t n, const int64_t dim, const int64_t* ids) {
 }
 
 DatasetPtr
-GenResultDataset(const void* tensor) {
+GenResultDataset(const int64_t rows, const int64_t dim, const void* tensor) {
     auto ret_ds = std::make_shared<Dataset>();
+    SetDatasetRows(ret_ds, rows);
+    SetDatasetDim(ret_ds, dim);
     SetDatasetOutputTensor(ret_ds, tensor);
     return ret_ds;
 }

diff --git a/knowhere/index/vector_index/adapter/VectorAdapter.h b/knowhere/index/vector_index/adapter/VectorAdapter.h
@@ -64,7 +64,6 @@ DEFINE_DATASET_SETTER(SetDatasetJsonIdSet, meta::JSON_ID_SET, const std::string)
 
 #define GET_DATA_WITH_IDS(ds_ptr)                     \
     auto rows = knowhere::GetDatasetRows(ds_ptr);     \
-    auto dim = knowhere::GetDatasetDim(ds_ptr);       \
     auto p_ids = knowhere::GetDatasetInputIDs(ds_ptr);
 
 #define GET_TENSOR_DATA(ds_ptr)                       \
@@ -82,7 +81,7 @@ extern DatasetPtr
 GenDatasetWithIds(const int64_t n, const int64_t dim, const int64_t* ids);
 
 extern DatasetPtr
-GenResultDataset(const void* tensor);
+GenResultDataset(const int64_t rows, const int64_t dim, const void* tensor);
 
 extern DatasetPtr
 GenResultDataset(const int64_t* ids, const float* distance);

diff --git a/knowhere/index/vector_offset_index/IndexIVF_NM.cpp b/knowhere/index/vector_offset_index/IndexIVF_NM.cpp
@@ -153,6 +153,7 @@ IVF_NM::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
     }
 
     GET_DATA_WITH_IDS(dataset_ptr)
+    auto dim = Dim();
 
     float* p_x = nullptr;
     auto release_when_exception = [&]() {
@@ -170,7 +171,7 @@ IVF_NM::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
             KNOWHERE_THROW_IF_NOT_FMT(id >= 0 && id < ivf_index->ntotal, "invalid id %ld", id);
             ivf_index->reconstruct_without_codes(id, p_x + i * dim);
         }
-        return GenResultDataset(p_x);
+        return GenResultDataset(rows, dim, p_x);
     } catch (faiss::FaissException& e) {
         release_when_exception();
         KNOWHERE_THROW_MSG(e.what());

diff --git a/knowhere/index/vector_offset_index/IndexIVF_NM.h b/knowhere/index/vector_offset_index/IndexIVF_NM.h
@@ -54,7 +54,7 @@ class IVF_NM : public VecIndex, public OffsetBaseIndex {
 
     bool
     HasRawData(const std::string& /*metric_type*/) const override {
-        return true;
+        return false;
     }
 
     DatasetPtr

diff --git a/unittest/test_annoy.cpp b/unittest/test_annoy.cpp
@@ -57,7 +57,7 @@ TEST_P(AnnoyTest, annoy_basic) {
     ASSERT_EQ(index_->Dim(), dim);
     ASSERT_GT(index_->Size(), 0);
 
-    ASSERT_TRUE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
+    ASSERT_FALSE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
     auto result = index_->GetVectorById(id_dataset, conf_);
     AssertVec(result, base_dataset, id_dataset, nq, dim);
 

diff --git a/unittest/test_binaryidmap.cpp b/unittest/test_binaryidmap.cpp
@@ -71,7 +71,7 @@ TEST_P(BinaryIDMAPTest, binaryidmap_basic) {
     EXPECT_EQ(index_->Dim(), dim);
     ASSERT_GT(index_->Size(), 0);
 
-    ASSERT_TRUE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
+    ASSERT_FALSE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
     auto result = index_->GetVectorById(id_dataset, conf_);
     AssertBinVec(result, base_dataset, id_dataset, nq, dim);
 

diff --git a/unittest/test_binaryivf.cpp b/unittest/test_binaryivf.cpp
@@ -68,7 +68,7 @@ TEST_P(BinaryIVFTest, binaryivf_basic) {
     EXPECT_EQ(index_->Dim(), dim);
     ASSERT_GT(index_->Size(), 0);
 
-    ASSERT_TRUE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
+    ASSERT_FALSE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
     auto result = index_->GetVectorById(id_dataset, conf_);
     AssertBinVec(result, base_dataset, id_dataset, nq, dim);
 

diff --git a/unittest/test_idmap.cpp b/unittest/test_idmap.cpp
@@ -90,7 +90,7 @@ TEST_P(IDMAPTest, idmap_basic) {
     EXPECT_EQ(index_->Dim(), dim);
     ASSERT_GT(index_->Size(), 0);
 
-    ASSERT_TRUE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
+    ASSERT_FALSE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
     auto result = index_->GetVectorById(id_dataset, conf_);
     AssertVec(result, base_dataset, id_dataset, nq, dim);
 

diff --git a/unittest/test_ivf_nm.cpp b/unittest/test_ivf_nm.cpp
@@ -105,7 +105,7 @@ TEST_P(IVFNMTest, ivfnm_basic) {
     LoadRawData(index_, base_dataset, conf_);
 
     if (index_mode_ == knowhere::IndexMode::MODE_CPU) {
-        ASSERT_TRUE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
+        ASSERT_FALSE(index_->HasRawData(knowhere::GetMetaMetricType(conf_)));
         auto result = index_->GetVectorById(id_dataset, conf_);
         AssertVec(result, base_dataset, id_dataset, nq, dim);