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);