diff --git a/core/src/db/engine/ExecutionEngine.h b/core/src/db/engine/ExecutionEngine.h index 86a014cf66ce2058b331372a0c41bd84acd2a766..c8784e8a90199c1684eb86a9f9babe27acfa36d3 100644 --- a/core/src/db/engine/ExecutionEngine.h +++ b/core/src/db/engine/ExecutionEngine.h @@ -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 { diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp index 3a2b940cec01f20a1eb16ffc6b474b18da4c813d..e2db7179bee15fc46866df8b223be96428921634 100644 --- a/core/src/db/engine/ExecutionEngineImpl.cpp +++ b/core/src/db/engine/ExecutionEngineImpl.cpp @@ -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; diff --git a/core/src/index/knowhere/CMakeLists.txt b/core/src/index/knowhere/CMakeLists.txt index 5f8e4d6970af5bf161af7258e40ac955e7b1ea2a..285461bdef65f89f32082598cc0b930207a4ba39 100644 --- a/core/src/index/knowhere/CMakeLists.txt +++ b/core/src/index/knowhere/CMakeLists.txt @@ -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 diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.cpp deleted file mode 100644 index 1bd45075e3f0f87bc3423391a5813eb61ac572b0..0000000000000000000000000000000000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.cpp +++ /dev/null @@ -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 -#include -#include -#include -#include - -#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 index_blobs; - std::vector index_len; - - // TODO(zirui): dev - // index_ptr_->SaveIndexToMemory(index_blobs, index_len); - BinarySet binary_set; - - // - // auto sample = std::make_shared(); - // sample.reset(static_cast(index_blobs[0])); - // auto tree = std::make_shared(); - // tree.reset(static_cast(index_blobs[1])); - // auto graph = std::make_shared(); - // graph.reset(static_cast(index_blobs[2])); - // auto metadata = std::make_shared(); - // metadata.reset(static_cast(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 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(); -//} - -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 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 diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.h deleted file mode 100644 index f6d436995b0d1051bf8c75c2c7a51cae4eb1d890..0000000000000000000000000000000000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexKDT.h +++ /dev/null @@ -1,87 +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. - -#pragma once - -#include -#include -#include -#include "VectorIndex.h" -#include "knowhere/index/IndexModel.h" - -namespace knowhere { - -class CPUKDTRNG : public VectorIndex { - public: - CPUKDTRNG() { - index_ptr_ = SPTAG::VectorIndex::CreateInstance(SPTAG::IndexAlgoType::KDT, SPTAG::VectorValueType::Float); - index_ptr_->SetParameter("DistCalcMethod", "L2"); - } - - 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; - - IndexModelPtr - Train(const DatasetPtr& dataset, const Config& config) override; - - void - Add(const DatasetPtr& dataset, const Config& config) override; - - DatasetPtr - Search(const DatasetPtr& dataset, const Config& config) override; - void - Seal() override; - - private: - void - SetParameters(const Config& config); - - private: - PreprocessorPtr preprocessor_; - std::shared_ptr index_ptr_; -}; - -using CPUKDTRNGPtr = std::shared_ptr; - -class CPUKDTRNGIndexModel : public IndexModel { - public: - BinarySet - Serialize() override; - - void - Load(const BinarySet& binary) override; - - private: - std::shared_ptr index_; -}; - -using CPUKDTRNGIndexModelPtr = std::shared_ptr; - -} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c64baebb93456a21c6773b792ceead78a554ba17 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.cpp @@ -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 +#include +#include + +#include +#include +#include + +#undef mkdir + +#include "knowhere/index/vector_index/IndexSPTAG.h" +#include "knowhere/index/vector_index/helpers/Definitions.h" +#include "knowhere/adapter/SptagAdapter.h" +#include "knowhere/common/Exception.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 index_blobs; + + std::shared_ptr> buffersize = index_ptr_->CalculateBufferSize(); + std::vector 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(); + sample.reset(static_cast(index_blobs[0].Data())); + auto tree = std::make_shared(); + tree.reset(static_cast(index_blobs[1].Data())); + auto graph = std::make_shared(); + graph.reset(static_cast(index_blobs[2].Data())); + auto deleteid = std::make_shared(); + deleteid.reset(static_cast(index_blobs[3].Data())); + auto metadata1 = std::make_shared(); + metadata1.reset(static_cast(index_blobs[4].Data())); + auto metadata2 = std::make_shared(); + metadata2.reset(static_cast(index_blobs[5].Data())); + auto config = std::make_shared(); + config.reset(static_cast((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(); +// 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(); +// data.reset(writer.data_); +// BinarySet binary_set; +// binary_set.Append("sptag", data, writer.total); +// auto graph = std::make_shared(); +// graph.reset(static_cast(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 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(config->data.get()); + +// std::vector 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 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(); +//} + + 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(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(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 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 diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h new file mode 100644 index 0000000000000000000000000000000000000000..c78807a3ee0d843abb2c522b1c939bcf8377b4e4 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexSPTAG.h @@ -0,0 +1,93 @@ +// 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. + +#pragma once + +#include + +#include +#include +#include + +#include "VectorIndex.h" +#include "knowhere/index/IndexModel.h" + +namespace knowhere { + + class CPUSPTAGRNG : public VectorIndex { + public: + 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; + + IndexModelPtr + Train(const DatasetPtr& dataset, const Config& config) override; + + void + Add(const DatasetPtr& dataset, const Config& config) override; + + DatasetPtr + Search(const DatasetPtr& dataset, const Config& config) override; + + void + Seal() override; + + private: + void + SetParameters(const Config& config); + + private: + PreprocessorPtr preprocessor_; + std::shared_ptr index_ptr_; + SPTAG::IndexAlgoType index_type_; + }; + + using CPUSPTAGRNGPtr = std::shared_ptr; + + class CPUSPTAGRNGIndexModel : public IndexModel { + public: + BinarySet + Serialize() override; + + void + Load(const BinarySet& binary) override; + + private: + std::shared_ptr index_; + }; + + using CPUSPTAGRNGIndexModelPtr = std::shared_ptr; + +} // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h index b2854abef8611edfb017873a17386c006fdf591b..9bb3cac303b9fcec5f91a6aabc1dd776ed0ce7a1 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/IndexParameter.h @@ -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; @@ -126,8 +152,57 @@ struct NSGCfg : public IVFCfg { }; using NSGConfig = std::shared_ptr; -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; + +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; + +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; } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp deleted file mode 100644 index 19bf070dbae735dd5142c823cf824f382f64c670..0000000000000000000000000000000000000000 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp +++ /dev/null @@ -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 - -#include "knowhere/index/vector_index/helpers/KDTParameterMgr.h" - -namespace knowhere { - -const std::vector& -KDTParameterMgr::GetKDTParameters() { - return kdt_parameters_; -} - -KDTParameterMgr::KDTParameterMgr() { - kdt_parameters_ = std::vector{ - {"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 diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp new file mode 100644 index 0000000000000000000000000000000000000000..f6e1409abdfce9fd38e02518261fb968c22c19a1 --- /dev/null +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.cpp @@ -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 + +#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(); + 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(); + 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 \ No newline at end of file diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h similarity index 55% rename from core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h rename to core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h index fe90761e173194e87e5e28cdc921e8f93df48107..d215d9e2750b746b6ad6f081caffc911842ebace 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/SPTAGParameterMgr.h @@ -22,31 +22,40 @@ #include #include +#include +#include "IndexParameter.h" + namespace knowhere { -using KDTParameter = std::pair; + using KDTConfig = std::shared_ptr; + using BKTConfig = std::shared_ptr; + + class SPTAGParameterMgr { + public: + const KDTConfig& + GetKDTParameters(); + + const BKTConfig& + GetBKTParameters(); -class KDTParameterMgr { - public: - const std::vector& - GetKDTParameters(); + public: + static SPTAGParameterMgr& + GetInstance() { + static SPTAGParameterMgr instance; + return instance; + } - public: - static KDTParameterMgr& - GetInstance() { - static KDTParameterMgr instance; - return instance; - } + SPTAGParameterMgr(const SPTAGParameterMgr&) = delete; - KDTParameterMgr(const KDTParameterMgr&) = delete; - KDTParameterMgr& - operator=(const KDTParameterMgr&) = delete; + SPTAGParameterMgr& + operator=(const SPTAGParameterMgr&) = delete; - private: - KDTParameterMgr(); + private: + SPTAGParameterMgr(); - private: - std::vector kdt_parameters_; -}; + private: + KDTConfig kdt_config_; + BKTConfig bkt_config_; + }; } // namespace knowhere diff --git a/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h b/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h index d00ea45365f4a4339599a7b0d861dc35104c13e5..0208f6d98357c2579ec22b429e66bd869942f14b 100644 --- a/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h +++ b/core/src/index/thirdparty/SPTAG/AnnService/inc/Core/Common/Dataset.h @@ -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; } diff --git a/core/src/index/unittest/CMakeLists.txt b/core/src/index/unittest/CMakeLists.txt index d3ba84e557e4f055d0571bdea3da40c691e407e0..93ae63a9ec65e1a88568dd383ebd12bee14760c3 100644 --- a/core/src/index/unittest/CMakeLists.txt +++ b/core/src/index/unittest/CMakeLists.txt @@ -82,17 +82,17 @@ if (NOT TARGET test_idmap) endif () target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs}) -# -set(kdt_srcs +# +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) diff --git a/core/src/index/unittest/test_kdt.cpp b/core/src/index/unittest/test_kdt.cpp deleted file mode 100644 index bbc7dcf94c903c59e3fa060198d16ed4814f4cb4..0000000000000000000000000000000000000000 --- a/core/src/index/unittest/test_kdt.cpp +++ /dev/null @@ -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 - -#include -#include - -#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(); - - auto tempconf = std::make_shared(); - tempconf->tptnubmber = 1; - tempconf->k = 10; - conf = tempconf; - - Init_with_default(); - } - - protected: - knowhere::Config conf; - std::shared_ptr 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(1, i * k + j) << " "; - // ss_dist << *dists->data()->GetValues(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(); -// 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 filename_list; -// std::vector> 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(); -// data.reset(load_data); -// load_data_list.Append(meta_list[i].first, data, bin_size); -// } -// -// auto new_index = std::make_shared(); -// new_index->Load(load_data_list); -// auto result = new_index->Search(query_dataset, conf); -// AssertAnns(result, nq, k); -// PrintResult(result, nq, k); -// } -//} diff --git a/core/src/index/unittest/test_sptag.cpp b/core/src/index/unittest/test_sptag.cpp new file mode 100644 index 0000000000000000000000000000000000000000..a0898a619ba01fe38704bcf1b021946651e8e3fe --- /dev/null +++ b/core/src/index/unittest/test_sptag.cpp @@ -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 + +#include +#include + +#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 { + protected: + void + SetUp() override { + IndexType = GetParam(); + Generate(128, 100, 5); + index_ = std::make_shared(IndexType); + if (IndexType == "KDT") { + auto tempconf = std::make_shared(); + tempconf->tptnumber = 1; + tempconf->k = 10; + conf = tempconf; + } else { + auto tempconf = std::make_shared(); + tempconf->tptnumber = 1; + tempconf->k = 10; + conf = tempconf; + } + + Init_with_default(); + } + + protected: + knowhere::Config conf; + std::shared_ptr 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(1, i * k + j) << " "; + //ss_dist << *dists->data()->GetValues(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(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 filename_list; + std::vector> 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(); + data.reset(load_data); + load_data_list.Append(meta_list[i].first, data, bin_size); + } + + auto new_index = std::make_shared(IndexType); + new_index->Load(load_data_list); + auto result = new_index->Search(query_dataset, conf); + AssertAnns(result, nq, k); + PrintResult(result, nq, k); + } +} diff --git a/core/src/index/unittest/utils.cpp b/core/src/index/unittest/utils.cpp index 2556b60fad6ce0d854e6dc34ab148e1092ed3757..110d8fbf922f83a00364d651d76eb12aedc7958c 100644 --- a/core/src/index/unittest/utils.cpp +++ b/core/src/index/unittest/utils.cpp @@ -153,22 +153,24 @@ void AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) { auto ids = result->ids(); for (auto i = 0; i < nq; i++) { - EXPECT_EQ(i, *((int64_t*)(ids) + i * k)); + EXPECT_EQ(i, *((int64_t*)(ids) + i * k)); // EXPECT_EQ(i, *(ids->data()->GetValues(1, i * 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(1, i * k + j)) << " "; - ss_dist << *(dists->data()->GetValues(1, i * k + j)) << " "; + //ss_id << *(ids->data()->GetValues(1, i * k + j)) << " "; + //ss_dist << *(dists->data()->GetValues(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; diff --git a/core/src/wrapper/ConfAdapter.cpp b/core/src/wrapper/ConfAdapter.cpp index 461745f1fdbaeb5386eb5fa89b4b8c0e942d64bf..4114eb2af40ce173f43562519410f18ca0b3257c 100644 --- a/core/src/wrapper/ConfAdapter.cpp +++ b/core/src/wrapper/ConfAdapter.cpp @@ -204,5 +204,35 @@ NSGConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) return conf; } +knowhere::Config +SPTAGKDTConfAdapter::Match(const TempMetaConf& metaconf) { + auto conf = std::make_shared(); + 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(); + conf->k = metaconf.k; + return conf; +} + +knowhere::Config +SPTAGBKTConfAdapter::Match(const TempMetaConf& metaconf) { + auto conf = std::make_shared(); + 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(); + conf->k = metaconf.k; + return conf; +} + } // namespace engine } // namespace milvus diff --git a/core/src/wrapper/ConfAdapter.h b/core/src/wrapper/ConfAdapter.h index 5ec3d524860e95dfdc0770442acda1ac6bde860d..600dd3a517aed1d5a8ea20011c21e979f1ca7fb0 100644 --- a/core/src/wrapper/ConfAdapter.h +++ b/core/src/wrapper/ConfAdapter.h @@ -97,5 +97,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 diff --git a/core/src/wrapper/ConfAdapterMgr.cpp b/core/src/wrapper/ConfAdapterMgr.cpp index d0eba0452934200c2e34298e326c6bb5849287e6..cf58c0110f12e0e9376bc4fa15c0733c95164040 100644 --- a/core/src/wrapper/ConfAdapterMgr.cpp +++ b/core/src/wrapper/ConfAdapterMgr.cpp @@ -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 diff --git a/core/src/wrapper/VecIndex.cpp b/core/src/wrapper/VecIndex.cpp index 81a13c60f90a4709d020f5c3d4bd664ffebb641e..9f68cac0a1bd189460ce2453bdbadd8e22256c45 100644 --- a/core/src/wrapper/VecIndex.cpp +++ b/core/src/wrapper/VecIndex.cpp @@ -22,7 +22,7 @@ #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/IndexSPTAG.h" #include "knowhere/index/vector_index/IndexNSG.h" #include "utils/Log.h" @@ -128,7 +128,11 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) { break; } case IndexType::SPTAG_KDT_RNT_CPU: { - index = std::make_shared(); + index = std::make_shared("KDT"); + break; + } + case IndexType::SPTAG_BKT_RNT_CPU: { + index = std::make_shared("BKT"); break; } case IndexType::FAISS_IVFSQ8_CPU: { diff --git a/core/src/wrapper/VecIndex.h b/core/src/wrapper/VecIndex.h index efe01a25d7b40b5209649e18bffe5102e4cc917c..e69655b087797cdc4f2794e9d79d0ffad167842d 100644 --- a/core/src/wrapper/VecIndex.h +++ b/core/src/wrapper/VecIndex.h @@ -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()); diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/unittest/wrapper/test_wrapper.cpp index 025601a1cd65d6d90cc8cb8aa92361d6e41bacef..413db60e37688c989cde742c26665b8e6715f195 100644 --- a/core/unittest/wrapper/test_wrapper.cpp +++ b/core/unittest/wrapper/test_wrapper.cpp @@ -29,15 +29,16 @@ INITIALIZE_EASYLOGGINGPP +using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; -using ::testing::Combine; class KnowhereWrapperTest - : public DataGenBase, - public TestWithParam<::std::tuple> { + : public DataGenBase, + public TestWithParam<::std::tuple> { protected: - void SetUp() override { + void + SetUp() override { #ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); #endif @@ -57,12 +58,13 @@ class KnowhereWrapperTest conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf); searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf); -// conf->k = k; -// conf->d = dim; -// conf->gpu_id = DEVICEID; + // conf->k = k; + // conf->d = dim; + // conf->gpu_id = DEVICEID; } - void TearDown() override { + void + TearDown() override { #ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); #endif @@ -75,22 +77,21 @@ 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)));