milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #438 from tinkerlin/0.6.0-#227 

0.6.0 #227
pull/441/head
Tinkerrr 2019-11-20 17:17:43 +08:00 committed by GitHub
parent 24e05a87d0 863cc5dbaf
commit e01eb29fff
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
23 changed files with 794 additions and 434 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@ -23,6 +23,7 @@ Please mark all change in change log and use the ticket from JIRA.
- \#77 - Support table partition
- \#127 - Support new Index type IVFPQ
- \#226 - Experimental shards middleware for Milvus
- \#227 - Support new index types SPTAG-KDT and SPTAG-BKT
- \#346 - Support build index with multiple gpu
## Improvement

4
core/src/db/engine/ExecutionEngine.h
View File

@ -35,7 +35,9 @@ enum class EngineType {
NSG_MIX,
FAISS_IVFSQ8H,
FAISS_PQ,
MAX_VALUE = FAISS_PQ,
SPTAG_KDT,
SPTAG_BKT,
MAX_VALUE = SPTAG_BKT,
};
enum class MetricType {

8
core/src/db/engine/ExecutionEngineImpl.cpp
View File

@ -124,6 +124,14 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) {
#endif
break;
}
case EngineType::SPTAG_KDT: {
index = GetVecIndexFactory(IndexType::SPTAG_KDT_RNT_CPU);
break;
}
case EngineType::SPTAG_BKT: {
index = GetVecIndexFactory(IndexType::SPTAG_BKT_RNT_CPU);
break;
}
default: {
ENGINE_LOG_ERROR << "Unsupported index type";
return nullptr;

4
core/src/index/knowhere/CMakeLists.txt
View File

@ -30,10 +30,10 @@ set(external_srcs
set(index_srcs
knowhere/index/preprocessor/Normalize.cpp
knowhere/index/vector_index/IndexKDT.cpp
knowhere/index/vector_index/IndexSPTAG.cpp
knowhere/index/vector_index/IndexIDMAP.cpp
knowhere/index/vector_index/IndexIVF.cpp
knowhere/index/vector_index/helpers/KDTParameterMgr.cpp
knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp
knowhere/index/vector_index/IndexNSG.cpp
knowhere/index/vector_index/nsg/NSG.cpp
knowhere/index/vector_index/nsg/NSGIO.cpp

180
core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.cpp
View File

@ -1,180 +0,0 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include <SPTAG/AnnService/inc/Core/Common.h>
#include <SPTAG/AnnService/inc/Core/VectorSet.h>
#include <SPTAG/AnnService/inc/Server/QueryParser.h>
#include <sstream>
#include <vector>
#undef mkdir
#include "knowhere/index/vector_index/IndexKDT.h"
#include "knowhere/index/vector_index/helpers/Definitions.h"
//#include "knowhere/index/preprocessor/normalize.h"
#include "knowhere/adapter/SptagAdapter.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/helpers/KDTParameterMgr.h"
namespace knowhere {
BinarySet
CPUKDTRNG::Serialize() {
std::vector<void*> index_blobs;
std::vector<int64_t> index_len;
// TODO(zirui): dev
// index_ptr_->SaveIndexToMemory(index_blobs, index_len);
BinarySet binary_set;
//
// auto sample = std::make_shared<uint8_t>();
// sample.reset(static_cast<uint8_t*>(index_blobs[0]));
// auto tree = std::make_shared<uint8_t>();
// tree.reset(static_cast<uint8_t*>(index_blobs[1]));
// auto graph = std::make_shared<uint8_t>();
// graph.reset(static_cast<uint8_t*>(index_blobs[2]));
// auto metadata = std::make_shared<uint8_t>();
// metadata.reset(static_cast<uint8_t*>(index_blobs[3]));
//
// binary_set.Append("samples", sample, index_len[0]);
// binary_set.Append("tree", tree, index_len[1]);
// binary_set.Append("graph", graph, index_len[2]);
// binary_set.Append("metadata", metadata, index_len[3]);
return binary_set;
}
void
CPUKDTRNG::Load(const BinarySet& binary_set) {
// TODO(zirui): dev
// std::vector<void*> index_blobs;
//
// auto samples = binary_set.GetByName("samples");
// index_blobs.push_back(samples->data.get());
//
// auto tree = binary_set.GetByName("tree");
// index_blobs.push_back(tree->data.get());
//
// auto graph = binary_set.GetByName("graph");
// index_blobs.push_back(graph->data.get());
//
// auto metadata = binary_set.GetByName("metadata");
// index_blobs.push_back(metadata->data.get());
//
// index_ptr_->LoadIndexFromMemory(index_blobs);
}
// PreprocessorPtr
// CPUKDTRNG::BuildPreprocessor(const DatasetPtr &dataset, const Config &config) {
// return std::make_shared<NormalizePreprocessor>();
//}
IndexModelPtr
CPUKDTRNG::Train(const DatasetPtr& origin, const Config& train_config) {
SetParameters(train_config);
DatasetPtr dataset = origin->Clone();
// if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
// && preprocessor_) {
// preprocessor_->Preprocess(dataset);
//}
auto vectorset = ConvertToVectorSet(dataset);
auto metaset = ConvertToMetadataSet(dataset);
index_ptr_->BuildIndex(vectorset, metaset);
// TODO: return IndexModelPtr
return nullptr;
}
void
CPUKDTRNG::Add(const DatasetPtr& origin, const Config& add_config) {
SetParameters(add_config);
DatasetPtr dataset = origin->Clone();
// if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
// && preprocessor_) {
// preprocessor_->Preprocess(dataset);
//}
auto vectorset = ConvertToVectorSet(dataset);
auto metaset = ConvertToMetadataSet(dataset);
index_ptr_->AddIndex(vectorset, metaset);
}
void
CPUKDTRNG::SetParameters(const Config& config) {
for (auto& para : KDTParameterMgr::GetInstance().GetKDTParameters()) {
// auto value = config.get_with_default(para.first, para.second);
index_ptr_->SetParameter(para.first, para.second);
}
}
DatasetPtr
CPUKDTRNG::Search(const DatasetPtr& dataset, const Config& config) {
SetParameters(config);
auto tensor = dataset->tensor()[0];
auto p = (float*)tensor->raw_mutable_data();
for (auto i = 0; i < 10; ++i) {
for (auto j = 0; j < 10; ++j) {
std::cout << p[i * 10 + j] << " ";
}
std::cout << std::endl;
}
std::vector<SPTAG::QueryResult> query_results = ConvertToQueryResult(dataset, config);
#pragma omp parallel for
for (auto i = 0; i < query_results.size(); ++i) {
auto target = (float*)query_results[i].GetTarget();
std::cout << target[0] << ", " << target[1] << ", " << target[2] << std::endl;
index_ptr_->SearchIndex(query_results[i]);
}
return ConvertToDataset(query_results);
}
int64_t
CPUKDTRNG::Count() {
index_ptr_->GetNumSamples();
}
int64_t
CPUKDTRNG::Dimension() {
index_ptr_->GetFeatureDim();
}
VectorIndexPtr
CPUKDTRNG::Clone() {
KNOWHERE_THROW_MSG("not support");
}
void
CPUKDTRNG::Seal() {
// do nothing
}
// TODO(linxj):
BinarySet
CPUKDTRNGIndexModel::Serialize() {
}
void
CPUKDTRNGIndexModel::Load(const BinarySet& binary) {
}
} // namespace knowhere

348
core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp Normal file
View File

@ -0,0 +1,348 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include <SPTAG/AnnService/inc/Core/Common.h>
#include <SPTAG/AnnService/inc/Core/VectorSet.h>
#include <SPTAG/AnnService/inc/Server/QueryParser.h>
#include <array>
#include <sstream>
#include <vector>
#undef mkdir
#include "knowhere/adapter/SptagAdapter.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexSPTAG.h"
#include "knowhere/index/vector_index/helpers/Definitions.h"
#include "knowhere/index/vector_index/helpers/SPTAGParameterMgr.h"
namespace knowhere {
CPUSPTAGRNG::CPUSPTAGRNG(const std::string& IndexType) {
if (IndexType == "KDT") {
index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::KDT, SPTAG::VectorValueType::Float);
index_ptr_->SetParameter("DistCalcMethod", "L2");
index_type_ = SPTAG::IndexAlgoType::KDT;
} else {
index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::BKT, SPTAG::VectorValueType::Float);
index_ptr_->SetParameter("DistCalcMethod", "L2");
index_type_ = SPTAG::IndexAlgoType::BKT;
}
}
BinarySet
CPUSPTAGRNG::Serialize() {
std::string index_config;
std::vector<SPTAG::ByteArray> index_blobs;
std::shared_ptr<std::vector<std::uint64_t>> buffersize = index_ptr_->CalculateBufferSize();
std::vector<char*> res(buffersize->size() + 1);
for (uint64_t i = 1; i < res.size(); i++) {
res[i] = new char[buffersize->at(i - 1)];
auto ptr = &res[i][0];
index_blobs.emplace_back(SPTAG::ByteArray((std::uint8_t*)ptr, buffersize->at(i - 1), false));
}
index_ptr_->SaveIndex(index_config, index_blobs);
size_t length = index_config.length();
char* cstr = new char[length];
snprintf(cstr, length, "%s", index_config.c_str());
BinarySet binary_set;
auto sample = std::make_shared<uint8_t>();
sample.reset(static_cast<uint8_t*>(index_blobs[0].Data()));
auto tree = std::make_shared<uint8_t>();
tree.reset(static_cast<uint8_t*>(index_blobs[1].Data()));
auto graph = std::make_shared<uint8_t>();
graph.reset(static_cast<uint8_t*>(index_blobs[2].Data()));
auto deleteid = std::make_shared<uint8_t>();
deleteid.reset(static_cast<uint8_t*>(index_blobs[3].Data()));
auto metadata1 = std::make_shared<uint8_t>();
metadata1.reset(static_cast<uint8_t*>(index_blobs[4].Data()));
auto metadata2 = std::make_shared<uint8_t>();
metadata2.reset(static_cast<uint8_t*>(index_blobs[5].Data()));
auto config = std::make_shared<uint8_t>();
config.reset(static_cast<uint8_t*>((void*)cstr));
binary_set.Append("samples", sample, index_blobs[0].Length());
binary_set.Append("tree", tree, index_blobs[1].Length());
binary_set.Append("deleteid", deleteid, index_blobs[3].Length());
binary_set.Append("metadata1", metadata1, index_blobs[4].Length());
binary_set.Append("metadata2", metadata2, index_blobs[5].Length());
binary_set.Append("config", config, length);
binary_set.Append("graph", graph, index_blobs[2].Length());
// MemoryIOWriter writer;
// size_t len = 0;
// for (int i = 0; i < 6; ++i) {
// len = index_blobs[i].Length();
// assert(len != 0);
// writer(&len, sizeof(size_t), 1);
// writer(index_blobs[i].Data(), len, 1);
// len = 0;
// }
// writer(&length, sizeof(size_t), 1);
// writer(cstr, length, 1);
// auto data = std::make_shared<uint8_t>();
// data.reset(writer.data_);
// BinarySet binary_set;
// binary_set.Append("sptag", data, writer.total);
// MemoryIOWriter writer;
// size_t len = 0;
// for (int i = 0; i < 6; ++i) {
// if (i == 2) continue;
// len = index_blobs[i].Length();
// assert(len != 0);
// writer(&len, sizeof(size_t), 1);
// writer(index_blobs[i].Data(), len, 1);
// len = 0;
// }
// writer(&length, sizeof(size_t), 1);
// writer(cstr, length, 1);
// auto data = std::make_shared<uint8_t>();
// data.reset(writer.data_);
// BinarySet binary_set;
// binary_set.Append("sptag", data, writer.total);
// auto graph = std::make_shared<uint8_t>();
// graph.reset(static_cast<uint8_t*>(index_blobs[2].Data()));
// binary_set.Append("graph", graph, index_blobs[2].Length());
return binary_set;
}
void
CPUSPTAGRNG::Load(const BinarySet& binary_set) {
std::string index_config;
std::vector<SPTAG::ByteArray> index_blobs;
auto samples = binary_set.GetByName("samples");
index_blobs.push_back(SPTAG::ByteArray(samples->data.get(), samples->size, false));
auto tree = binary_set.GetByName("tree");
index_blobs.push_back(SPTAG::ByteArray(tree->data.get(), tree->size, false));
auto graph = binary_set.GetByName("graph");
index_blobs.push_back(SPTAG::ByteArray(graph->data.get(), graph->size, false));
auto deleteid = binary_set.GetByName("deleteid");
index_blobs.push_back(SPTAG::ByteArray(deleteid->data.get(), deleteid->size, false));
auto metadata1 = binary_set.GetByName("metadata1");
index_blobs.push_back(SPTAG::ByteArray(metadata1->data.get(), metadata1->size, false));
auto metadata2 = binary_set.GetByName("metadata2");
index_blobs.push_back(SPTAG::ByteArray(metadata2->data.get(), metadata2->size, false));
auto config = binary_set.GetByName("config");
index_config = reinterpret_cast<char*>(config->data.get());
// std::vector<SPTAG::ByteArray> index_blobs;
// auto data = binary_set.GetByName("sptag");
// MemoryIOReader reader;
// reader.total = data->size;
// reader.data_ = data->data.get();
// size_t len = 0;
// for (int i = 0; i < 6; ++i) {
// reader(&len, sizeof(size_t), 1);
// assert(len != 0);
// auto binary = new uint8_t[len];
// reader(binary, len, 1);
// index_blobs.emplace_back(SPTAG::ByteArray(binary, len, true));
// len = 0;
// }
// reader(&len, sizeof(size_t), 1);
// assert(len != 0);
// auto config = new char[len];
// reader(config, len, 1);
// std::string index_config = config;
// delete[] config;
// std::vector<SPTAG::ByteArray> index_blobs;
// auto data = binary_set.GetByName("sptag");
// MemoryIOReader reader;
// reader.total = data->size;
// reader.data_ = data->data.get();
// size_t len = 0;
// for (int i = 0; i < 6; ++i) {
// if (i == 2) {
// auto graph = binary_set.GetByName("graph");
// index_blobs.emplace_back(SPTAG::ByteArray(graph->data.get(), graph->size, false));
// continue;
// }
// reader(&len, sizeof(size_t), 1);
// assert(len != 0);
// auto binary = new uint8_t[len];
// reader(binary, len, 1);
// index_blobs.emplace_back(SPTAG::ByteArray(binary, len, true));
// len = 0;
// }
// reader(&len, sizeof(size_t), 1);
// assert(len != 0);
// auto config = new char[len];
// reader(config, len, 1);
// std::string index_config = config;
// delete[] config;
index_ptr_->LoadIndex(index_config, index_blobs);
}
// PreprocessorPtr
// CPUKDTRNG::BuildPreprocessor(const DatasetPtr &dataset, const Config &config) {
// return std::make_shared<NormalizePreprocessor>();
//}
IndexModelPtr
CPUSPTAGRNG::Train(const DatasetPtr& origin, const Config& train_config) {
SetParameters(train_config);
DatasetPtr dataset = origin->Clone();
// if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
// && preprocessor_) {
// preprocessor_->Preprocess(dataset);
//}
auto vectorset = ConvertToVectorSet(dataset);
auto metaset = ConvertToMetadataSet(dataset);
index_ptr_->BuildIndex(vectorset, metaset);
// TODO: return IndexModelPtr
return nullptr;
}
void
CPUSPTAGRNG::Add(const DatasetPtr& origin, const Config& add_config) {
// SetParameters(add_config);
// DatasetPtr dataset = origin->Clone();
//
// // if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
// // && preprocessor_) {
// // preprocessor_->Preprocess(dataset);
// //}
//
// auto vectorset = ConvertToVectorSet(dataset);
// auto metaset = ConvertToMetadataSet(dataset);
// index_ptr_->AddIndex(vectorset, metaset);
}
void
CPUSPTAGRNG::SetParameters(const Config& config) {
#define Assign(param_name, str_name) \
conf->param_name == INVALID_VALUE ? index_ptr_->SetParameter(str_name, std::to_string(build_cfg->param_name)) \
: index_ptr_->SetParameter(str_name, std::to_string(conf->param_name))
if (index_type_ == SPTAG::IndexAlgoType::KDT) {
auto conf = std::dynamic_pointer_cast<KDTCfg>(config);
auto build_cfg = SPTAGParameterMgr::GetInstance().GetKDTParameters();
Assign(kdtnumber, "KDTNumber");
Assign(numtopdimensionkdtsplit, "NumTopDimensionKDTSplit");
Assign(samples, "Samples");
Assign(tptnumber, "TPTNumber");
Assign(tptleafsize, "TPTLeafSize");
Assign(numtopdimensiontptsplit, "NumTopDimensionTPTSplit");
Assign(neighborhoodsize, "NeighborhoodSize");
Assign(graphneighborhoodscale, "GraphNeighborhoodScale");
Assign(graphcefscale, "GraphCEFScale");
Assign(refineiterations, "RefineIterations");
Assign(cef, "CEF");
Assign(maxcheckforrefinegraph, "MaxCheckForRefineGraph");
Assign(numofthreads, "NumberOfThreads");
Assign(maxcheck, "MaxCheck");
Assign(thresholdofnumberofcontinuousnobetterpropagation, "ThresholdOfNumberOfContinuousNoBetterPropagation");
Assign(numberofinitialdynamicpivots, "NumberOfInitialDynamicPivots");
Assign(numberofotherdynamicpivots, "NumberOfOtherDynamicPivots");
} else {
auto conf = std::dynamic_pointer_cast<BKTCfg>(config);
auto build_cfg = SPTAGParameterMgr::GetInstance().GetBKTParameters();
Assign(bktnumber, "BKTNumber");
Assign(bktkmeansk, "BKTKMeansK");
Assign(bktleafsize, "BKTLeafSize");
Assign(samples, "Samples");
Assign(tptnumber, "TPTNumber");
Assign(tptleafsize, "TPTLeafSize");
Assign(numtopdimensiontptsplit, "NumTopDimensionTPTSplit");
Assign(neighborhoodsize, "NeighborhoodSize");
Assign(graphneighborhoodscale, "GraphNeighborhoodScale");
Assign(graphcefscale, "GraphCEFScale");
Assign(refineiterations, "RefineIterations");
Assign(cef, "CEF");
Assign(maxcheckforrefinegraph, "MaxCheckForRefineGraph");
Assign(numofthreads, "NumberOfThreads");
Assign(maxcheck, "MaxCheck");
Assign(thresholdofnumberofcontinuousnobetterpropagation, "ThresholdOfNumberOfContinuousNoBetterPropagation");
Assign(numberofinitialdynamicpivots, "NumberOfInitialDynamicPivots");
Assign(numberofotherdynamicpivots, "NumberOfOtherDynamicPivots");
}
}
DatasetPtr
CPUSPTAGRNG::Search(const DatasetPtr& dataset, const Config& config) {
SetParameters(config);
auto tensor = dataset->tensor()[0];
auto p = (float*)tensor->raw_mutable_data();
for (auto i = 0; i < 10; ++i) {
for (auto j = 0; j < 10; ++j) {
std::cout << p[i * 10 + j] << " ";
}
std::cout << std::endl;
}
std::vector<SPTAG::QueryResult> query_results = ConvertToQueryResult(dataset, config);
#pragma omp parallel for
for (auto i = 0; i < query_results.size(); ++i) {
auto target = (float*)query_results[i].GetTarget();
std::cout << target[0] << ", " << target[1] << ", " << target[2] << std::endl;
index_ptr_->SearchIndex(query_results[i]);
}
return ConvertToDataset(query_results);
}
int64_t
CPUSPTAGRNG::Count() {
return index_ptr_->GetNumSamples();
}
int64_t
CPUSPTAGRNG::Dimension() {
return index_ptr_->GetFeatureDim();
}
VectorIndexPtr
CPUSPTAGRNG::Clone() {
KNOWHERE_THROW_MSG("not support");
}
void
CPUSPTAGRNG::Seal() {
return; // do nothing
}
BinarySet
CPUSPTAGRNGIndexModel::Serialize() {
// KNOWHERE_THROW_MSG("not support"); // not support
}
void
CPUSPTAGRNGIndexModel::Load(const BinarySet& binary) {
// KNOWHERE_THROW_MSG("not support"); // not support
}
} // namespace knowhere

22
core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.h → core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h
View File

@ -18,33 +18,37 @@
#pragma once
#include <SPTAG/AnnService/inc/Core/VectorIndex.h>
#include <cstdint>
#include <memory>
#include <string>
#include "VectorIndex.h"
#include "knowhere/index/IndexModel.h"
namespace knowhere {
class CPUKDTRNG : public VectorIndex {
class CPUSPTAGRNG : public VectorIndex {
public:
CPUKDTRNG() {
index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::KDT, SPTAG::VectorValueType::Float);
index_ptr_->SetParameter("DistCalcMethod", "L2");
}
explicit CPUSPTAGRNG(const std::string& IndexType);
public:
BinarySet
Serialize() override;
VectorIndexPtr
Clone() override;
void
Load(const BinarySet& index_array) override;
public:
// PreprocessorPtr
// BuildPreprocessor(const DatasetPtr &dataset, const Config &config) override;
int64_t
Count() override;
int64_t
Dimension() override;
@ -56,6 +60,7 @@ class CPUKDTRNG : public VectorIndex {
DatasetPtr
Search(const DatasetPtr& dataset, const Config& config) override;
void
Seal() override;
@ -66,11 +71,12 @@ class CPUKDTRNG : public VectorIndex {
private:
PreprocessorPtr preprocessor_;
std::shared_ptr<SPTAG::VectorIndex> index_ptr_;
SPTAG::IndexAlgoType index_type_;
};
using CPUKDTRNGPtr = std::shared_ptr<CPUKDTRNG>;
using CPUSPTAGRNGPtr = std::shared_ptr<CPUSPTAGRNG>;
class CPUKDTRNGIndexModel : public IndexModel {
class CPUSPTAGRNGIndexModel : public IndexModel {
public:
BinarySet
Serialize() override;
@ -82,6 +88,6 @@ class CPUKDTRNGIndexModel : public IndexModel {
std::shared_ptr<SPTAG::VectorIndex> index_;
};
using CPUKDTRNGIndexModelPtr = std::shared_ptr<CPUKDTRNGIndexModel>;
using CPUSPTAGRNGIndexModelPtr = std::shared_ptr<CPUSPTAGRNGIndexModel>;
} // namespace knowhere

79
core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h
View File

@ -42,6 +42,32 @@ constexpr int64_t DEFAULT_OUT_DEGREE = INVALID_VALUE;
constexpr int64_t DEFAULT_CANDIDATE_SISE = INVALID_VALUE;
constexpr int64_t DEFAULT_NNG_K = INVALID_VALUE;
// SPTAG Config
constexpr int64_t DEFAULT_SAMPLES = INVALID_VALUE;
constexpr int64_t DEFAULT_TPTNUMBER = INVALID_VALUE;
constexpr int64_t DEFAULT_TPTLEAFSIZE = INVALID_VALUE;
constexpr int64_t DEFAULT_NUMTOPDIMENSIONTPTSPLIT = INVALID_VALUE;
constexpr int64_t DEFAULT_NEIGHBORHOODSIZE = INVALID_VALUE;
constexpr int64_t DEFAULT_GRAPHNEIGHBORHOODSCALE = INVALID_VALUE;
constexpr int64_t DEFAULT_GRAPHCEFSCALE = INVALID_VALUE;
constexpr int64_t DEFAULT_REFINEITERATIONS = INVALID_VALUE;
constexpr int64_t DEFAULT_CEF = INVALID_VALUE;
constexpr int64_t DEFAULT_MAXCHECKFORREFINEGRAPH = INVALID_VALUE;
constexpr int64_t DEFAULT_NUMOFTHREADS = INVALID_VALUE;
constexpr int64_t DEFAULT_MAXCHECK = INVALID_VALUE;
constexpr int64_t DEFAULT_THRESHOLDOFNUMBEROFCONTINUOUSNOBETTERPROPAGATION = INVALID_VALUE;
constexpr int64_t DEFAULT_NUMBEROFINITIALDYNAMICPIVOTS = INVALID_VALUE;
constexpr int64_t DEFAULT_NUMBEROFOTHERDYNAMICPIVOTS = INVALID_VALUE;
// KDT Config
constexpr int64_t DEFAULT_KDTNUMBER = INVALID_VALUE;
constexpr int64_t DEFAULT_NUMTOPDIMENSIONKDTSPLIT = INVALID_VALUE;
// BKT Config
constexpr int64_t DEFAULT_BKTNUMBER = INVALID_VALUE;
constexpr int64_t DEFAULT_BKTKMEANSK = INVALID_VALUE;
constexpr int64_t DEFAULT_BKTLEAFSIZE = INVALID_VALUE;
struct IVFCfg : public Cfg {
int64_t nlist = DEFAULT_NLIST;
int64_t nprobe = DEFAULT_NPROBE;
@ -135,8 +161,57 @@ struct NSGCfg : public IVFCfg {
};
using NSGConfig = std::shared_ptr<NSGCfg>;
struct KDTCfg : public Cfg {
int64_t tptnubmber = -1;
struct SPTAGCfg : public Cfg {
int64_t samples = DEFAULT_SAMPLES;
int64_t tptnumber = DEFAULT_TPTNUMBER;
int64_t tptleafsize = DEFAULT_TPTLEAFSIZE;
int64_t numtopdimensiontptsplit = DEFAULT_NUMTOPDIMENSIONTPTSPLIT;
int64_t neighborhoodsize = DEFAULT_NEIGHBORHOODSIZE;
int64_t graphneighborhoodscale = DEFAULT_GRAPHNEIGHBORHOODSCALE;
int64_t graphcefscale = DEFAULT_GRAPHCEFSCALE;
int64_t refineiterations = DEFAULT_REFINEITERATIONS;
int64_t cef = DEFAULT_CEF;
int64_t maxcheckforrefinegraph = DEFAULT_MAXCHECKFORREFINEGRAPH;
int64_t numofthreads = DEFAULT_NUMOFTHREADS;
int64_t maxcheck = DEFAULT_MAXCHECK;
int64_t thresholdofnumberofcontinuousnobetterpropagation = DEFAULT_THRESHOLDOFNUMBEROFCONTINUOUSNOBETTERPROPAGATION;
int64_t numberofinitialdynamicpivots = DEFAULT_NUMBEROFINITIALDYNAMICPIVOTS;
int64_t numberofotherdynamicpivots = DEFAULT_NUMBEROFOTHERDYNAMICPIVOTS;
SPTAGCfg() = default;
bool
CheckValid() override {
return true;
};
};
using SPTAGConfig = std::shared_ptr<SPTAGCfg>;
struct KDTCfg : public SPTAGCfg {
int64_t kdtnumber = DEFAULT_KDTNUMBER;
int64_t numtopdimensionkdtsplit = DEFAULT_NUMTOPDIMENSIONKDTSPLIT;
KDTCfg() = default;
bool
CheckValid() override {
return true;
};
};
using KDTConfig = std::shared_ptr<KDTCfg>;
struct BKTCfg : public SPTAGCfg {
int64_t bktnumber = DEFAULT_BKTNUMBER;
int64_t bktkmeansk = DEFAULT_BKTKMEANSK;
int64_t bktleafsize = DEFAULT_BKTLEAFSIZE;
BKTCfg() = default;
bool
CheckValid() override {
return true;
};
};
using BKTConfig = std::shared_ptr<BKTCfg>;
} // namespace knowhere

55
core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp
View File

@ -1,55 +0,0 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include <mutex>
#include "knowhere/index/vector_index/helpers/KDTParameterMgr.h"
namespace knowhere {
const std::vector<KDTParameter>&
KDTParameterMgr::GetKDTParameters() {
return kdt_parameters_;
}
KDTParameterMgr::KDTParameterMgr() {
kdt_parameters_ = std::vector<KDTParameter>{
{"KDTNumber", "1"},
{"NumTopDimensionKDTSplit", "5"},
{"NumSamplesKDTSplitConsideration", "100"},
{"TPTNumber", "1"},
{"TPTLeafSize", "2000"},
{"NumTopDimensionTPTSplit", "5"},
{"NeighborhoodSize", "32"},
{"GraphNeighborhoodScale", "2"},
{"GraphCEFScale", "2"},
{"RefineIterations", "0"},
{"CEF", "1000"},
{"MaxCheckForRefineGraph", "10000"},
{"NumberOfThreads", "1"},
{"MaxCheck", "8192"},
{"ThresholdOfNumberOfContinuousNoBetterPropagation", "3"},
{"NumberOfInitialDynamicPivots", "50"},
{"NumberOfOtherDynamicPivots", "4"},
};
}
} // namespace knowhere

75
core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp Normal file
View File

@ -0,0 +1,75 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include <mutex>
#include "knowhere/index/vector_index/helpers/SPTAGParameterMgr.h"
namespace knowhere {
const KDTConfig&
SPTAGParameterMgr::GetKDTParameters() {
return kdt_config_;
}
const BKTConfig&
SPTAGParameterMgr::GetBKTParameters() {
return bkt_config_;
}
SPTAGParameterMgr::SPTAGParameterMgr() {
kdt_config_ = std::make_shared<KDTCfg>();
kdt_config_->kdtnumber = 1;
kdt_config_->numtopdimensionkdtsplit = 5;
kdt_config_->samples = 100;
kdt_config_->tptnumber = 1;
kdt_config_->tptleafsize = 2000;
kdt_config_->numtopdimensiontptsplit = 5;
kdt_config_->neighborhoodsize = 32;
kdt_config_->graphneighborhoodscale = 2;
kdt_config_->graphcefscale = 2;
kdt_config_->refineiterations = 0;
kdt_config_->cef = 1000;
kdt_config_->maxcheckforrefinegraph = 10000;
kdt_config_->numofthreads = 1;
kdt_config_->maxcheck = 8192;
kdt_config_->thresholdofnumberofcontinuousnobetterpropagation = 3;
kdt_config_->numberofinitialdynamicpivots = 50;
kdt_config_->numberofotherdynamicpivots = 4;
bkt_config_ = std::make_shared<BKTCfg>();
bkt_config_->bktnumber = 1;
bkt_config_->bktkmeansk = 32;
bkt_config_->bktleafsize = 8;
bkt_config_->samples = 100;
bkt_config_->tptnumber = 1;
bkt_config_->tptleafsize = 2000;
bkt_config_->numtopdimensiontptsplit = 5;
bkt_config_->neighborhoodsize = 32;
bkt_config_->graphneighborhoodscale = 2;
bkt_config_->graphcefscale = 2;
bkt_config_->refineiterations = 0;
bkt_config_->cef = 1000;
bkt_config_->maxcheckforrefinegraph = 10000;
bkt_config_->numofthreads = 1;
bkt_config_->maxcheck = 8192;
bkt_config_->thresholdofnumberofcontinuousnobetterpropagation = 3;
bkt_config_->numberofinitialdynamicpivots = 50;
bkt_config_->numberofotherdynamicpivots = 4;
}
} // namespace knowhere

29
core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h → core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h
View File

@ -22,31 +22,40 @@
#include <utility>
#include <vector>
#include <SPTAG/AnnService/inc/Core/Common.h>
#include "IndexParameter.h"
namespace knowhere {
using KDTParameter = std::pair<std::string, std::string>;
using KDTConfig = std::shared_ptr<KDTCfg>;
using BKTConfig = std::shared_ptr<BKTCfg>;
class KDTParameterMgr {
class SPTAGParameterMgr {
public:
const std::vector<KDTParameter>&
const KDTConfig&
GetKDTParameters();
const BKTConfig&
GetBKTParameters();
public:
static KDTParameterMgr&
static SPTAGParameterMgr&
GetInstance() {
static KDTParameterMgr instance;
static SPTAGParameterMgr instance;
return instance;
}
KDTParameterMgr(const KDTParameterMgr&) = delete;
KDTParameterMgr&
operator=(const KDTParameterMgr&) = delete;
SPTAGParameterMgr(const SPTAGParameterMgr&) = delete;
SPTAGParameterMgr&
operator=(const SPTAGParameterMgr&) = delete;
private:
KDTParameterMgr();
SPTAGParameterMgr();
private:
std::vector<KDTParameter> kdt_parameters_;
KDTConfig kdt_config_;
BKTConfig bkt_config_;
};
} // namespace knowhere

2
core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h vendored
View File

@ -195,7 +195,7 @@ namespace SPTAG
C = *((DimensionType*)pDataPointsMemFile);
pDataPointsMemFile += sizeof(DimensionType);
Initialize(R, C, (T*)pDataPointsMemFile);
Initialize(R, C, (T*)pDataPointsMemFile, false);
std::cout << "Load " << name << " (" << R << ", " << C << ") Finish!" << std::endl;
return true;
}

16
core/src/index/unittest/CMakeLists.txt
View File

@ -82,17 +82,17 @@ if (NOT TARGET test_idmap)
endif ()
target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs})
#<KDT-TEST>
set(kdt_srcs
#<SPTAG-TEST>
set(sptag_srcs
${INDEX_SOURCE_DIR}/knowhere/knowhere/adapter/SptagAdapter.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/preprocessor/Normalize.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexKDT.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp
)
if (NOT TARGET test_kdt)
add_executable(test_kdt test_kdt.cpp ${kdt_srcs} ${util_srcs})
if (NOT TARGET test_sptag)
add_executable(test_sptag test_sptag.cpp ${sptag_srcs} ${util_srcs})
endif ()
target_link_libraries(test_kdt
target_link_libraries(test_sptag
SPTAGLibStatic
${depend_libs} ${unittest_libs} ${basic_libs})
@ -106,7 +106,7 @@ endif ()
install(TARGETS test_ivf DESTINATION unittest)
install(TARGETS test_idmap DESTINATION unittest)
install(TARGETS test_kdt DESTINATION unittest)
install(TARGETS test_sptag DESTINATION unittest)
if (KNOWHERE_GPU_VERSION)
install(TARGETS test_gpuresource DESTINATION unittest)
install(TARGETS test_customized_index DESTINATION unittest)

144
core/src/index/unittest/test_kdt.cpp
View File

@ -1,144 +0,0 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include <iostream>
#include <sstream>
#include "knowhere/adapter/SptagAdapter.h"
#include "knowhere/adapter/Structure.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexKDT.h"
#include "knowhere/index/vector_index/helpers/Definitions.h"
#include "unittest/utils.h"
using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
class KDTTest : public DataGen, public ::testing::Test {
protected:
void
SetUp() override {
Generate(96, 1000, 10);
index_ = std::make_shared<knowhere::CPUKDTRNG>();
auto tempconf = std::make_shared<knowhere::KDTCfg>();
tempconf->tptnubmber = 1;
tempconf->k = 10;
conf = tempconf;
Init_with_default();
}
protected:
knowhere::Config conf;
std::shared_ptr<knowhere::CPUKDTRNG> index_ = nullptr;
};
// TODO(lxj): add test about count() and dimension()
TEST_F(KDTTest, kdt_basic) {
assert(!xb.empty());
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
auto result = index_->Search(query_dataset, conf);
AssertAnns(result, nq, k);
{
// auto ids = result->array()[0];
// auto dists = result->array()[1];
auto ids = result->ids();
auto dists = result->dist();
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < nq; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *((int64_t*)(ids) + i * k + j) << " ";
ss_dist << *((float*)(dists) + i * k + j) << " ";
// ss_id << *ids->data()->GetValues<int64_t>(1, i * k + j) << " ";
// ss_dist << *dists->data()->GetValues<float>(1, i * k + j) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
}
// TODO(zirui): enable test
// TEST_F(KDTTest, kdt_serialize) {
// assert(!xb.empty());
//
// auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
// index_->set_preprocessor(preprocessor);
//
// auto model = index_->Train(base_dataset, conf);
// // index_->Add(base_dataset, conf);
// auto binaryset = index_->Serialize();
// auto new_index = std::make_shared<knowhere::CPUKDTRNG>();
// new_index->Load(binaryset);
// auto result = new_index->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
// PrintResult(result, nq, k);
// ASSERT_EQ(new_index->Count(), nb);
// ASSERT_EQ(new_index->Dimension(), dim);
// ASSERT_THROW({ new_index->Clone(); }, knowhere::KnowhereException);
// ASSERT_NO_THROW({ new_index->Seal(); });
//
// {
// int fileno = 0;
// const std::string& base_name = "/tmp/kdt_serialize_test_bin_";
// std::vector<std::string> filename_list;
// std::vector<std::pair<std::string, size_t>> meta_list;
// for (auto& iter : binaryset.binary_map_) {
// const std::string& filename = base_name + std::to_string(fileno);
// FileIOWriter writer(filename);
// writer(iter.second->data.get(), iter.second->size);
//
// meta_list.emplace_back(std::make_pair(iter.first, iter.second->size));
// filename_list.push_back(filename);
// ++fileno;
// }
//
// knowhere::BinarySet load_data_list;
// for (int i = 0; i < filename_list.size() && i < meta_list.size(); ++i) {
// auto bin_size = meta_list[i].second;
// FileIOReader reader(filename_list[i]);
//
// auto load_data = new uint8_t[bin_size];
// reader(load_data, bin_size);
// auto data = std::make_shared<uint8_t>();
// data.reset(load_data);
// load_data_list.Append(meta_list[i].first, data, bin_size);
// }
//
// auto new_index = std::make_shared<knowhere::CPUKDTRNG>();
// new_index->Load(load_data_list);
// auto result = new_index->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
// PrintResult(result, nq, k);
// }
//}

154
core/src/index/unittest/test_sptag.cpp Normal file
View File

@ -0,0 +1,154 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include <iostream>
#include <sstream>
#include "knowhere/adapter/SptagAdapter.h"
#include "knowhere/adapter/Structure.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexSPTAG.h"
#include "knowhere/index/vector_index/helpers/Definitions.h"
#include "unittest/utils.h"
using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
class SPTAGTest : public DataGen, public TestWithParam<std::string> {
protected:
void
SetUp() override {
IndexType = GetParam();
Generate(128, 100, 5);
index_ = std::make_shared<knowhere::CPUSPTAGRNG>(IndexType);
if (IndexType == "KDT") {
auto tempconf = std::make_shared<knowhere::KDTCfg>();
tempconf->tptnumber = 1;
tempconf->k = 10;
conf = tempconf;
} else {
auto tempconf = std::make_shared<knowhere::BKTCfg>();
tempconf->tptnumber = 1;
tempconf->k = 10;
conf = tempconf;
}
Init_with_default();
}
protected:
knowhere::Config conf;
std::shared_ptr<knowhere::CPUSPTAGRNG> index_ = nullptr;
std::string IndexType;
};
INSTANTIATE_TEST_CASE_P(SPTAGParameters, SPTAGTest, Values("KDT", "BKT"));
// TODO(lxj): add test about count() and dimension()
TEST_P(SPTAGTest, sptag_basic) {
assert(!xb.empty());
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
auto result = index_->Search(query_dataset, conf);
AssertAnns(result, nq, k);
{
// auto ids = result->array()[0];
// auto dists = result->array()[1];
auto ids = result->ids();
auto dists = result->dist();
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < nq; i++) {
for (auto j = 0; j < k; ++j) {
// ss_id << *ids->data()->GetValues<int64_t>(1, i * k + j) << " ";
// ss_dist << *dists->data()->GetValues<float>(1, i * k + j) << " ";
ss_id << *((int64_t*)(ids) + i * k + j) << " ";
ss_dist << *((float*)(dists) + i * k + j) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
}
TEST_P(SPTAGTest, sptag_serialize) {
assert(!xb.empty());
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->Add(base_dataset, conf);
auto binaryset = index_->Serialize();
auto new_index = std::make_shared<knowhere::CPUSPTAGRNG>(IndexType);
new_index->Load(binaryset);
auto result = new_index->Search(query_dataset, conf);
AssertAnns(result, nq, k);
PrintResult(result, nq, k);
ASSERT_EQ(new_index->Count(), nb);
ASSERT_EQ(new_index->Dimension(), dim);
// ASSERT_THROW({ new_index->Clone(); }, knowhere::KnowhereException);
// ASSERT_NO_THROW({ new_index->Seal(); });
{
int fileno = 0;
const std::string& base_name = "/tmp/sptag_serialize_test_bin_";
std::vector<std::string> filename_list;
std::vector<std::pair<std::string, size_t>> meta_list;
for (auto& iter : binaryset.binary_map_) {
const std::string& filename = base_name + std::to_string(fileno);
FileIOWriter writer(filename);
writer(iter.second->data.get(), iter.second->size);
meta_list.emplace_back(std::make_pair(iter.first, iter.second->size));
filename_list.push_back(filename);
++fileno;
}
knowhere::BinarySet load_data_list;
for (int i = 0; i < filename_list.size() && i < meta_list.size(); ++i) {
auto bin_size = meta_list[i].second;
FileIOReader reader(filename_list[i]);
auto load_data = new uint8_t[bin_size];
reader(load_data, bin_size);
auto data = std::make_shared<uint8_t>();
data.reset(load_data);
load_data_list.Append(meta_list[i].first, data, bin_size);
}
auto new_index = std::make_shared<knowhere::CPUSPTAGRNG>(IndexType);
new_index->Load(load_data_list);
auto result = new_index->Search(query_dataset, conf);
AssertAnns(result, nq, k);
PrintResult(result, nq, k);
}
}

12
core/src/index/unittest/utils.cpp
View File

@ -160,15 +160,17 @@ AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
auto ids = result->ids();
auto dists = result->dist();
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto i = 0; i < nq; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
// ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
// ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
ss_id << *((int64_t*)(ids) + i * k + j) << " ";
ss_dist << *((float*)(dists) + i * k + j) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;

30
core/src/wrapper/ConfAdapter.cpp
View File

@ -201,5 +201,35 @@ NSGConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type)
return conf;
}
knowhere::Config
SPTAGKDTConfAdapter::Match(const TempMetaConf& metaconf) {
auto conf = std::make_shared<knowhere::KDTCfg>();
conf->d = metaconf.dim;
conf->metric_type = metaconf.metric_type;
return conf;
}
knowhere::Config
SPTAGKDTConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) {
auto conf = std::make_shared<knowhere::KDTCfg>();
conf->k = metaconf.k;
return conf;
}
knowhere::Config
SPTAGBKTConfAdapter::Match(const TempMetaConf& metaconf) {
auto conf = std::make_shared<knowhere::BKTCfg>();
conf->d = metaconf.dim;
conf->metric_type = metaconf.metric_type;
return conf;
}
knowhere::Config
SPTAGBKTConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) {
auto conf = std::make_shared<knowhere::BKTCfg>();
conf->k = metaconf.k;
return conf;
}
} // namespace engine
} // namespace milvus

18
core/src/wrapper/ConfAdapter.h
View File

@ -94,5 +94,23 @@ class NSGConfAdapter : public IVFConfAdapter {
MatchSearch(const TempMetaConf& metaconf, const IndexType& type) final;
};
class SPTAGKDTConfAdapter : public ConfAdapter {
public:
knowhere::Config
Match(const TempMetaConf& metaconf) override;
knowhere::Config
MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override;
};
class SPTAGBKTConfAdapter : public ConfAdapter {
public:
knowhere::Config
Match(const TempMetaConf& metaconf) override;
knowhere::Config
MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override;
};
} // namespace engine
} // namespace milvus

3
core/src/wrapper/ConfAdapterMgr.cpp
View File

@ -56,6 +56,9 @@ AdapterMgr::RegisterAdapter() {
REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_MIX, ivfpq_mix);
REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexType::NSG_MIX, nsg_mix);
REGISTER_CONF_ADAPTER(SPTAGKDTConfAdapter, IndexType::SPTAG_KDT_RNT_CPU, sptag_kdt);
REGISTER_CONF_ADAPTER(SPTAGBKTConfAdapter, IndexType::SPTAG_BKT_RNT_CPU, sptag_bkt);
}
} // namespace engine

2
core/src/wrapper/VecImpl.cpp
View File

@ -26,8 +26,8 @@
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#endif

8
core/src/wrapper/VecIndex.cpp
View File

@ -22,8 +22,8 @@
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#include "knowhere/index/vector_index/IndexKDT.h"
#include "knowhere/index/vector_index/IndexNSG.h"
#include "knowhere/index/vector_index/IndexSPTAG.h"
#include "utils/Log.h"
#ifdef MILVUS_GPU_VERSION
@ -128,7 +128,11 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) {
break;
}
case IndexType::SPTAG_KDT_RNT_CPU: {
index = std::make_shared<knowhere::CPUKDTRNG>();
index = std::make_shared<knowhere::CPUSPTAGRNG>("KDT");
break;
}
case IndexType::SPTAG_BKT_RNT_CPU: {
index = std::make_shared<knowhere::CPUSPTAGRNG>("BKT");
break;
}
case IndexType::FAISS_IVFSQ8_CPU: {

4
core/src/wrapper/VecIndex.h
View File

@ -49,6 +49,7 @@ enum class IndexType {
FAISS_IVFSQ8_HYBRID, // only support build on gpu.
NSG_MIX,
FAISS_IVFPQ_MIX,
SPTAG_BKT_RNT_CPU,
};
class VecIndex;
@ -139,6 +140,9 @@ write_index(VecIndexPtr index, const std::string& location);
extern VecIndexPtr
read_index(const std::string& location);
VecIndexPtr
read_index(const std::string& location, knowhere::BinarySet& index_binary);
extern VecIndexPtr
GetVecIndexFactory(const IndexType& type, const Config& cfg = Config());

30
core/unittest/wrapper/test_wrapper.cpp
View File

@ -16,28 +16,29 @@
// under the License.
#include "easyloggingpp/easylogging++.h"
#include "wrapper/VecIndex.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#endif
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "wrapper/VecIndex.h"
#include "wrapper/utils.h"
#include <gtest/gtest.h>
INITIALIZE_EASYLOGGINGPP
using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
using ::testing::Combine;
class KnowhereWrapperTest
: public DataGenBase,
public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
: public DataGenBase,
public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
protected:
void SetUp() override {
void
SetUp() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
#endif
@ -58,7 +59,8 @@ class KnowhereWrapperTest
searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf);
}
void TearDown() override {
void
TearDown() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
@ -71,22 +73,20 @@ class KnowhereWrapperTest
knowhere::Config searchconf;
};
INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
Values(
//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"]
INSTANTIATE_TEST_CASE_P(
WrapperParam, KnowhereWrapperTest,
Values(
//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"]
#ifdef MILVUS_GPU_VERSION
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 1000, 10, 10),
// std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB,
// 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFPQ_MIX, "Default", 64, 1000, 10, 10),
// std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10),
// std::make_tuple(milvus::engine::IndexType::NSG_MIX, "Default", 128, 250000, 10, 10),
#endif
// std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10),
// std::make_tuple(milvus::engine::IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 100, 10, 10),
// std::make_tuple(milvus::engine::IndexType::SPTAG_BKT_RNT_CPU, "Default", 128, 100, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 1000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 1000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10)));