Skip to content
This repository has been archived by the owner on Aug 16, 2023. It is now read-only.

Commit

Permalink
Fix GetVectorById Dim, Return Dataset and Set All HasRawData to False…
Browse files Browse the repository at this point in the history
… Except HNSW (#957)

Signed-off-by: Patrick Weizhi Xu <[email protected]>
  • Loading branch information
PwzXxm authored Jun 26, 2023
1 parent b9bdbfc commit b6f01a6
Show file tree
Hide file tree
Showing 20 changed files with 29 additions and 21 deletions.
3 changes: 2 additions & 1 deletion knowhere/index/vector_index/IndexAnnoy.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion knowhere/index/vector_index/IndexAnnoy.h
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ class IndexAnnoy : public VecIndex {

bool
HasRawData(const std::string& /*metric_type*/) const override {
return true;
return false;
}

DatasetPtr
Expand Down
3 changes: 2 additions & 1 deletion knowhere/index/vector_index/IndexBinaryIDMAP.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [&]() {
Expand All @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion knowhere/index/vector_index/IndexBinaryIDMAP.h
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ class BinaryIDMAP : public VecIndex, public FaissBaseBinaryIndex {

bool
HasRawData(const std::string& /*metric_type*/) const override {
return true;
return false;
}

DatasetPtr
Expand Down
3 changes: 2 additions & 1 deletion knowhere/index/vector_index/IndexBinaryIVF.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [&]() {
Expand All @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion knowhere/index/vector_index/IndexBinaryIVF.h
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ class BinaryIVF : public VecIndex, public FaissBaseBinaryIndex {

bool
HasRawData(const std::string& /*metric_type*/) const override {
return true;
return false;
}

DatasetPtr
Expand Down
3 changes: 2 additions & 1 deletion knowhere/index/vector_index/IndexHNSW.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand Down
3 changes: 2 additions & 1 deletion knowhere/index/vector_index/IndexIDMAP.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [&]() {
Expand All @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion knowhere/index/vector_index/IndexIDMAP.h
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ class IDMAP : public VecIndex, public FaissBaseIndex {

bool
HasRawData(const std::string& /*metric_type*/) const override {
return true;
return false;
}

DatasetPtr
Expand Down
3 changes: 2 additions & 1 deletion knowhere/index/vector_index/IndexIVF.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [&]() {
Expand All @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion knowhere/index/vector_index/IndexIVF.h
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ class IVF : public VecIndex, public FaissBaseIndex {

bool
HasRawData(const std::string& /*metric_type*/) const override {
return true;
return false;
}

DatasetPtr
Expand Down
4 changes: 3 additions & 1 deletion knowhere/index/vector_index/adapter/VectorAdapter.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
3 changes: 1 addition & 2 deletions knowhere/index/vector_index/adapter/VectorAdapter.h
Original file line number Diff line number Diff line change
Expand Up @@ -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) \
Expand All @@ -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);
Expand Down
3 changes: 2 additions & 1 deletion knowhere/index/vector_offset_index/IndexIVF_NM.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [&]() {
Expand All @@ -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());
Expand Down
2 changes: 1 addition & 1 deletion knowhere/index/vector_offset_index/IndexIVF_NM.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion unittest/test_annoy.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
2 changes: 1 addition & 1 deletion unittest/test_binaryidmap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
2 changes: 1 addition & 1 deletion unittest/test_binaryivf.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
2 changes: 1 addition & 1 deletion unittest/test_idmap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
2 changes: 1 addition & 1 deletion unittest/test_ivf_nm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down

0 comments on commit b6f01a6

Please sign in to comment.