mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-29 23:15:28 +08:00
247 lines
10 KiB
C++
247 lines
10 KiB
C++
// Licensed to the LF AI & Data foundation 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 "index/VectorDiskIndex.h"
|
|
#include "index/Meta.h"
|
|
#include "index/Utils.h"
|
|
|
|
#include "storage/LocalChunkManager.h"
|
|
#include "config/ConfigKnowhere.h"
|
|
#include "storage/Util.h"
|
|
#include "common/Consts.h"
|
|
#include "common/Utils.h"
|
|
#include "common/RangeSearchHelper.h"
|
|
|
|
namespace milvus::index {
|
|
|
|
#ifdef BUILD_DISK_ANN
|
|
|
|
#define kSearchListMaxValue1 200 // used if tok <= 20
|
|
#define kSearchListMaxValue2 65535 // used for topk > 20
|
|
#define kPrepareDim 100
|
|
#define kPrepareRows 1
|
|
|
|
template <typename T>
|
|
VectorDiskAnnIndex<T>::VectorDiskAnnIndex(const IndexType& index_type,
|
|
const MetricType& metric_type,
|
|
const IndexMode& index_mode,
|
|
storage::FileManagerImplPtr file_manager)
|
|
: VectorIndex(index_type, index_mode, metric_type) {
|
|
file_manager_ = std::dynamic_pointer_cast<storage::DiskFileManagerImpl>(file_manager);
|
|
auto& local_chunk_manager = storage::LocalChunkManager::GetInstance();
|
|
auto local_index_path_prefix = file_manager_->GetLocalIndexObjectPrefix();
|
|
|
|
// As we have guarded dup-load in QueryNode,
|
|
// this assertion failed only if the Milvus rebooted in the same pod,
|
|
// need to remove these files then re-load the segment
|
|
if (local_chunk_manager.Exist(local_index_path_prefix)) {
|
|
local_chunk_manager.RemoveDir(local_index_path_prefix);
|
|
}
|
|
|
|
local_chunk_manager.CreateDir(local_index_path_prefix);
|
|
auto diskann_index_pack = knowhere::Pack(std::shared_ptr<knowhere::FileManager>(file_manager));
|
|
index_ = knowhere::IndexFactory::Instance().Create(GetIndexType(), diskann_index_pack);
|
|
}
|
|
|
|
template <typename T>
|
|
void
|
|
VectorDiskAnnIndex<T>::Load(const BinarySet& binary_set /* not used */, const Config& config) {
|
|
knowhere::Json load_config = update_load_json(config);
|
|
|
|
auto index_files = GetValueFromConfig<std::vector<std::string>>(config, "index_files");
|
|
AssertInfo(index_files.has_value(), "index file paths is empty when load disk ann index data");
|
|
file_manager_->CacheIndexToDisk(index_files.value());
|
|
|
|
// todo : replace by index::load function later
|
|
knowhere::DataSetPtr qs = std::make_unique<knowhere::DataSet>();
|
|
qs->SetRows(kPrepareRows);
|
|
qs->SetDim(kPrepareDim);
|
|
qs->SetIsOwner(true);
|
|
auto query = new T[kPrepareRows * kPrepareDim];
|
|
qs->SetTensor(query);
|
|
index_.Search(*qs, load_config, nullptr);
|
|
|
|
SetDim(index_.Dim());
|
|
}
|
|
|
|
template <typename T>
|
|
void
|
|
VectorDiskAnnIndex<T>::BuildWithDataset(const DatasetPtr& dataset, const Config& config) {
|
|
auto& local_chunk_manager = storage::LocalChunkManager::GetInstance();
|
|
knowhere::Json build_config;
|
|
build_config.update(config);
|
|
// set data path
|
|
auto segment_id = file_manager_->GetFileDataMeta().segment_id;
|
|
auto field_id = file_manager_->GetFileDataMeta().field_id;
|
|
auto local_data_path = storage::GenFieldRawDataPathPrefix(segment_id, field_id) + "raw_data";
|
|
build_config[DISK_ANN_RAW_DATA_PATH] = local_data_path;
|
|
|
|
auto local_index_path_prefix = file_manager_->GetLocalIndexObjectPrefix();
|
|
build_config[DISK_ANN_PREFIX_PATH] = local_index_path_prefix;
|
|
|
|
auto num_threads = GetValueFromConfig<std::string>(build_config, DISK_ANN_BUILD_THREAD_NUM);
|
|
AssertInfo(num_threads.has_value(), "param " + std::string(DISK_ANN_BUILD_THREAD_NUM) + "is empty");
|
|
build_config[DISK_ANN_THREADS_NUM] = std::atoi(num_threads.value().c_str());
|
|
|
|
if (!local_chunk_manager.Exist(local_data_path)) {
|
|
local_chunk_manager.CreateFile(local_data_path);
|
|
}
|
|
|
|
int64_t offset = 0;
|
|
auto num = uint32_t(milvus::GetDatasetRows(dataset));
|
|
local_chunk_manager.Write(local_data_path, offset, &num, sizeof(num));
|
|
offset += sizeof(num);
|
|
|
|
auto dim = uint32_t(milvus::GetDatasetDim(dataset));
|
|
local_chunk_manager.Write(local_data_path, offset, &dim, sizeof(dim));
|
|
offset += sizeof(dim);
|
|
|
|
auto data_size = num * dim * sizeof(float);
|
|
auto raw_data = const_cast<void*>(milvus::GetDatasetTensor(dataset));
|
|
local_chunk_manager.Write(local_data_path, offset, raw_data, data_size);
|
|
|
|
knowhere::DataSet* ds_ptr = nullptr;
|
|
index_.Build(*ds_ptr, build_config);
|
|
|
|
local_chunk_manager.RemoveDir(storage::GetSegmentRawDataPathPrefix(segment_id));
|
|
// TODO ::
|
|
// SetDim(index_->Dim());
|
|
}
|
|
|
|
template <typename T>
|
|
std::unique_ptr<SearchResult>
|
|
VectorDiskAnnIndex<T>::Query(const DatasetPtr dataset, const SearchInfo& search_info, const BitsetView& bitset) {
|
|
AssertInfo(GetMetricType() == search_info.metric_type_,
|
|
"Metric type of field index isn't the same with search info");
|
|
auto num_queries = dataset->GetRows();
|
|
auto topk = search_info.topk_;
|
|
|
|
knowhere::Json search_config = search_info.search_params_;
|
|
|
|
search_config[knowhere::meta::TOPK] = topk;
|
|
search_config[knowhere::meta::METRIC_TYPE] = GetMetricType();
|
|
|
|
// set search list size
|
|
auto search_list_size = GetValueFromConfig<uint32_t>(search_info.search_params_, DISK_ANN_QUERY_LIST);
|
|
AssertInfo(search_list_size.has_value(), "param " + std::string(DISK_ANN_QUERY_LIST) + "is empty");
|
|
AssertInfo(search_list_size.value() >= topk, "search_list should be greater than or equal to topk");
|
|
AssertInfo(search_list_size.value() <= std::max(uint32_t(topk * 10), uint32_t(kSearchListMaxValue1)) &&
|
|
search_list_size.value() <= uint32_t(kSearchListMaxValue2),
|
|
"search_list should be less than max(topk*10, 200) and less than 65535");
|
|
search_config[DISK_ANN_SEARCH_LIST_SIZE] = search_list_size.value();
|
|
|
|
// set beamwidth
|
|
search_config[DISK_ANN_QUERY_BEAMWIDTH] = int(search_beamwidth_);
|
|
|
|
// set index prefix, will be removed later
|
|
auto local_index_path_prefix = file_manager_->GetLocalIndexObjectPrefix();
|
|
search_config[DISK_ANN_PREFIX_PATH] = local_index_path_prefix;
|
|
|
|
// set json reset field, will be removed later
|
|
search_config[DISK_ANN_PQ_CODE_BUDGET] = 0.0;
|
|
|
|
auto final = [&] {
|
|
auto radius = GetValueFromConfig<float>(search_info.search_params_, RADIUS);
|
|
if (radius.has_value()) {
|
|
search_config[RADIUS] = radius.value();
|
|
auto range_filter = GetValueFromConfig<float>(search_info.search_params_, RANGE_FILTER);
|
|
if (range_filter.has_value()) {
|
|
search_config[RANGE_FILTER] = range_filter.value();
|
|
CheckRangeSearchParam(search_config[RADIUS], search_config[RANGE_FILTER], GetMetricType());
|
|
}
|
|
auto res = index_.RangeSearch(*dataset, search_config, bitset);
|
|
|
|
if (!res.has_value()) {
|
|
PanicCodeInfo(ErrorCodeEnum::UnexpectedError,
|
|
"failed to range search, " + MatchKnowhereError(res.error()));
|
|
}
|
|
return SortRangeSearchResult(res.value(), topk, num_queries, GetMetricType());
|
|
} else {
|
|
auto res = index_.Search(*dataset, search_config, bitset);
|
|
if (!res.has_value()) {
|
|
PanicCodeInfo(ErrorCodeEnum::UnexpectedError, "failed to search, " + MatchKnowhereError(res.error()));
|
|
}
|
|
return res.value();
|
|
}
|
|
}();
|
|
|
|
auto ids = final->GetIds();
|
|
float* distances = const_cast<float*>(final->GetDistance());
|
|
final->SetIsOwner(true);
|
|
|
|
auto round_decimal = search_info.round_decimal_;
|
|
auto total_num = num_queries * topk;
|
|
|
|
if (round_decimal != -1) {
|
|
const float multiplier = pow(10.0, round_decimal);
|
|
for (int i = 0; i < total_num; i++) {
|
|
distances[i] = round(distances[i] * multiplier) / multiplier;
|
|
}
|
|
}
|
|
auto result = std::make_unique<SearchResult>();
|
|
result->seg_offsets_.resize(total_num);
|
|
result->distances_.resize(total_num);
|
|
result->total_nq_ = num_queries;
|
|
result->unity_topK_ = topk;
|
|
|
|
std::copy_n(ids, total_num, result->seg_offsets_.data());
|
|
std::copy_n(distances, total_num, result->distances_.data());
|
|
|
|
return result;
|
|
}
|
|
|
|
template <typename T>
|
|
void
|
|
VectorDiskAnnIndex<T>::CleanLocalData() {
|
|
auto& local_chunk_manager = storage::LocalChunkManager::GetInstance();
|
|
local_chunk_manager.RemoveDir(file_manager_->GetLocalIndexObjectPrefix());
|
|
local_chunk_manager.RemoveDir(file_manager_->GetLocalRawDataObjectPrefix());
|
|
}
|
|
|
|
template <typename T>
|
|
inline knowhere::Json
|
|
VectorDiskAnnIndex<T>::update_load_json(const Config& config) {
|
|
knowhere::Json load_config;
|
|
load_config.update(config);
|
|
|
|
// set data path
|
|
auto local_index_path_prefix = file_manager_->GetLocalIndexObjectPrefix();
|
|
load_config[DISK_ANN_PREFIX_PATH] = local_index_path_prefix;
|
|
|
|
// set base info
|
|
load_config[DISK_ANN_PREPARE_WARM_UP] = false;
|
|
load_config[DISK_ANN_PREPARE_USE_BFS_CACHE] = false;
|
|
|
|
// set threads number
|
|
auto num_threads = GetValueFromConfig<std::string>(load_config, DISK_ANN_LOAD_THREAD_NUM);
|
|
AssertInfo(num_threads.has_value(), "param " + std::string(DISK_ANN_LOAD_THREAD_NUM) + "is empty");
|
|
load_config[DISK_ANN_THREADS_NUM] = std::atoi(num_threads.value().c_str());
|
|
|
|
// update search_beamwidth
|
|
auto beamwidth = GetValueFromConfig<std::string>(load_config, DISK_ANN_QUERY_BEAMWIDTH);
|
|
if (beamwidth.has_value()) {
|
|
search_beamwidth_ = std::atoi(beamwidth.value().c_str());
|
|
}
|
|
|
|
return load_config;
|
|
}
|
|
|
|
template class VectorDiskAnnIndex<float>;
|
|
|
|
#endif
|
|
|
|
} // namespace milvus::index
|