wehub-resource-sync
1788 lines
57 KiB
C++
1788 lines
57 KiB
C++
// Copyright 2025-present the zvec project
|
|
//
|
|
// Licensed 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 "helper.h"
|
|
|
|
mutex recall_lock;
|
|
bool g_compare_by_id = false;
|
|
float g_recall_precision;
|
|
|
|
|
|
//--------------------------------------------------
|
|
// Recall
|
|
//--------------------------------------------------
|
|
enum RetrievalMode { RM_UNDEFINED = 0, RM_DENSE = 1, RM_SPARSE = 2 };
|
|
|
|
enum FilterMode { FM_UNDEFINED = 0, FM_NONE = 1, FM_TAG = 2 };
|
|
|
|
template <typename T>
|
|
class Recall {
|
|
public:
|
|
Recall(size_t threads, const string &output, size_t batch_count,
|
|
FilterMode filter_mode)
|
|
: threads_(threads),
|
|
output_(output),
|
|
batch_count_(batch_count),
|
|
filter_mode_{filter_mode} {
|
|
if (threads_ == 0) {
|
|
pool_ = make_shared<ThreadPool>(true);
|
|
threads_ = pool_->count();
|
|
cout << "Using cpu count as thread pool count[" << threads_ << "]"
|
|
<< endl;
|
|
} else {
|
|
pool_ = make_shared<ThreadPool>(threads_, true);
|
|
cout << "Using thread pool count[" << threads_ << "]" << endl;
|
|
}
|
|
if (batch_count_ < 1) {
|
|
batch_count_ = 1;
|
|
call_batch_api_ = false;
|
|
} else {
|
|
call_batch_api_ = true;
|
|
}
|
|
}
|
|
|
|
static void stop(int signo) {
|
|
if (STOP_NOW) {
|
|
exit(signo);
|
|
}
|
|
STOP_NOW = true;
|
|
cout << "\rTrying to stop. press [Ctrl+C] again kill immediately." << endl
|
|
<< flush;
|
|
}
|
|
|
|
void run_dense(core_interface::Index::Pointer index,
|
|
core_interface::BaseIndexQueryParam::Pointer query_param,
|
|
const string &recall_tops, size_t gt_count) {
|
|
StringHelper::Split(recall_tops, ",", &topk_ids_);
|
|
std::sort(topk_ids_.begin(), topk_ids_.end());
|
|
|
|
for (auto i : topk_ids_) {
|
|
recall_res_[i] = 0.0f;
|
|
}
|
|
size_t topk = recall_res_.rbegin()->first;
|
|
|
|
gt_count = topk < gt_count ? gt_count : topk;
|
|
|
|
if (external_gt_file_enabled_) {
|
|
cout << "Internal ground truth file NOT used since external ground truth "
|
|
"file has been loaded"
|
|
<< endl;
|
|
} else {
|
|
cout << "Loading internal ground truth file" << endl;
|
|
|
|
if (!load_gt_dense(index, gt_count)) {
|
|
LOG_ERROR("Load ground truth file failed!");
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (batch_queries_.size() < threads_) {
|
|
threads_ = batch_queries_.size();
|
|
pool_ = make_shared<ThreadPool>(true, threads_);
|
|
cout << "Query size too small, resize thread pool count[" << threads_
|
|
<< "]" << endl;
|
|
}
|
|
|
|
// Prepare file handler
|
|
vector<pair<fstream *, fstream *>> output_fs;
|
|
if (!output_.empty()) {
|
|
if (!ailego::FileHelper::MakePath(output_.c_str())) {
|
|
LOG_ERROR("make path %s failed", output_.c_str());
|
|
return;
|
|
}
|
|
struct stat sb;
|
|
if (stat(output_.c_str(), &sb) == 0 && S_ISDIR(sb.st_mode)) {
|
|
cout << "logs output to : " << output_ << endl;
|
|
for (size_t i = 0; i < threads_; ++i) {
|
|
fstream *fs_k = new fstream();
|
|
fs_k->open(output_ + "/t" + to_string(i) + ".knn", ios::out);
|
|
fstream *fs_l = new fstream();
|
|
fs_l->open(output_ + "/t" + to_string(i) + ".linear", ios::out);
|
|
output_fs.push_back(make_pair(fs_k, fs_l));
|
|
}
|
|
}
|
|
}
|
|
|
|
signal(SIGINT, stop);
|
|
size_t i = 0;
|
|
for (; !STOP_NOW && i < batch_queries_.size();) {
|
|
if (pool_->pending_count() >= pool_->count()) {
|
|
this_thread::sleep_for(chrono::microseconds(1));
|
|
continue;
|
|
}
|
|
|
|
Closure::Pointer task =
|
|
Closure::New(this, &Recall::recall_one_dense, index, query_param,
|
|
topk, i, output_fs);
|
|
pool_->enqueue_and_wake(task);
|
|
|
|
i++;
|
|
}
|
|
pool_->wait_finish();
|
|
|
|
for (auto fs : output_fs) {
|
|
fs.first->close();
|
|
fs.second->close();
|
|
delete fs.first;
|
|
delete fs.second;
|
|
}
|
|
cout << "Process query: " << i << endl;
|
|
for (auto it : recall_res_) {
|
|
cout << "Recall@" << it.first << ": "
|
|
<< it.second / linear_queries_.size() << endl;
|
|
}
|
|
}
|
|
|
|
bool load_query(const std::string &query_file, const std::string &first_sep,
|
|
const std::string &second_sep) {
|
|
TxtInputReader<T> reader;
|
|
|
|
if (!reader.load_query(query_file, first_sep, second_sep, linear_queries_,
|
|
linear_sparse_data_, linear_taglists_)) {
|
|
LOG_ERROR("Load query error");
|
|
return false;
|
|
}
|
|
|
|
if (batch_count_ == 1) {
|
|
batch_queries_ = linear_queries_;
|
|
|
|
for (size_t i = 0; i < linear_sparse_data_.size(); ++i) {
|
|
vector<uint32_t> sparse_count;
|
|
sparse_count.push_back(linear_sparse_data_[i].count);
|
|
|
|
batch_sparse_counts_.push_back(sparse_count);
|
|
batch_sparse_indices_.push_back(linear_sparse_data_[i].indices);
|
|
batch_sparse_features_.push_back(linear_sparse_data_[i].features);
|
|
}
|
|
|
|
for (size_t i = 0; i < linear_taglists_.size(); ++i) {
|
|
vector<vector<uint64_t>> new_taglists;
|
|
new_taglists.push_back(linear_taglists_[i]);
|
|
|
|
batch_taglists_.push_back(std::move(new_taglists));
|
|
}
|
|
} else {
|
|
size_t num_batch =
|
|
(linear_queries_.size() + batch_count_ - 1) / batch_count_;
|
|
size_t idx = 0;
|
|
for (size_t n = 0; n < num_batch; ++n) {
|
|
vector<T> batch_query;
|
|
vector<uint32_t> batch_sparse_count;
|
|
vector<uint32_t> batch_sparse_indices;
|
|
vector<T> batch_sparse_feature;
|
|
vector<vector<uint64_t>> batch_taglists;
|
|
|
|
for (size_t i = 0; i < batch_count_; ++i) {
|
|
for (size_t k = 0; k < linear_queries_[idx].size(); ++k) {
|
|
batch_query.push_back(linear_queries_[idx][k]);
|
|
}
|
|
|
|
batch_sparse_count.push_back(linear_sparse_data_[idx].count);
|
|
|
|
for (size_t k = 0; k < linear_sparse_data_[idx].indices.size(); ++k) {
|
|
batch_sparse_indices.push_back(linear_sparse_data_[idx].indices[k]);
|
|
}
|
|
|
|
for (size_t k = 0; k < linear_sparse_data_[idx].features.size();
|
|
++k) {
|
|
batch_sparse_feature.push_back(
|
|
linear_sparse_data_[idx].features[k]);
|
|
}
|
|
|
|
idx = (idx + 1) % linear_queries_.size();
|
|
}
|
|
|
|
batch_queries_.push_back(batch_query);
|
|
batch_sparse_counts_.push_back(batch_sparse_count);
|
|
batch_sparse_indices_.push_back(batch_sparse_indices);
|
|
batch_sparse_features_.push_back(batch_sparse_feature);
|
|
batch_taglists_.push_back(batch_taglists);
|
|
}
|
|
}
|
|
|
|
dim_ = linear_queries_[0].size();
|
|
total_querys_ = linear_queries_.size();
|
|
if (typeid(T) == typeid(float)) {
|
|
qmeta_.set_meta(IndexMeta::DataType::DT_FP32, dim_);
|
|
} else if (typeid(T) == typeid(int8_t)) {
|
|
qmeta_.set_meta(IndexMeta::DataType::DT_INT8, dim_);
|
|
} else {
|
|
LOG_ERROR("unsupported type");
|
|
return false;
|
|
}
|
|
|
|
cout << "Load query done!" << endl;
|
|
return true;
|
|
}
|
|
|
|
bool load_external_gt_file(const std::string &external_gt_file,
|
|
const std::string &first_sep,
|
|
const std::string &second_sep) {
|
|
TxtInputReader<T> reader;
|
|
bool ret =
|
|
reader.load_external_gt(external_gt_file, first_sep, second_sep, gt_);
|
|
if (ret) {
|
|
cout << "Load external ground truth file["
|
|
<< File::BaseName(external_gt_file) << "] done!" << endl;
|
|
external_gt_file_enabled_ = true;
|
|
} else {
|
|
LOG_ERROR("Failed to load ground truth file!");
|
|
}
|
|
|
|
return ret;
|
|
}
|
|
|
|
private:
|
|
std::string compute_crc(size_t gt_count) {
|
|
uint32_t crc = 0u;
|
|
// dense
|
|
if (batch_queries_.size() > 0) {
|
|
size_t one_size = dim_ * sizeof(T);
|
|
size_t data_size = total_querys_ * one_size + sizeof(size_t);
|
|
char *data = new char[data_size];
|
|
size_t q = 0;
|
|
char *p = data;
|
|
for (; q < batch_queries_.size(); ++q) {
|
|
memcpy(p, batch_queries_[q].data(),
|
|
batch_queries_[q].size() * sizeof(T));
|
|
p += batch_queries_[q].size() * sizeof(T);
|
|
}
|
|
memcpy(p, >_count, sizeof(size_t));
|
|
crc = Crc32c::Hash(data, data_size, crc);
|
|
delete[] data;
|
|
}
|
|
|
|
// sparse
|
|
if (linear_sparse_data_.size() > 0) {
|
|
for (size_t i = 0; i < linear_sparse_data_.size(); ++i) {
|
|
crc = Crc32c::Hash(&(linear_sparse_data_[i].count), sizeof(uint32_t),
|
|
crc);
|
|
crc =
|
|
Crc32c::Hash(linear_sparse_data_[i].indices.data(),
|
|
linear_sparse_data_[i].count * sizeof(uint32_t), crc);
|
|
crc = Crc32c::Hash(linear_sparse_data_[i].features.data(),
|
|
linear_sparse_data_[i].count * sizeof(T), crc);
|
|
}
|
|
}
|
|
|
|
char crc_str[64];
|
|
snprintf(crc_str, sizeof(crc_str), "%X", crc);
|
|
|
|
return std::string(crc_str);
|
|
}
|
|
|
|
bool load_gt_dense(core_interface::Index::Pointer index, size_t gt_count) {
|
|
std::string crc_str = compute_crc(gt_count);
|
|
|
|
string gt_file = string("gt.") + crc_str;
|
|
|
|
File gtf;
|
|
if (!gtf.IsRegular(gt_file.c_str())) {
|
|
cout << "Ground truth file[" << gt_file << "] not exist, try to create it"
|
|
<< endl;
|
|
ElapsedTime timer;
|
|
|
|
size_t size = sizeof(uint64_t) + sizeof(float);
|
|
size_t file_size =
|
|
linear_queries_.size() * (sizeof(int) + size * gt_count);
|
|
|
|
std::string gt_file_temp = gt_file + ".tmp";
|
|
gtf.create(gt_file_temp.c_str(), file_size);
|
|
|
|
gt_.resize(linear_queries_.size());
|
|
|
|
atomic_bool error(false);
|
|
size_t count = 0;
|
|
float s = linear_queries_.size() / 100.0;
|
|
size_t pc = 0;
|
|
SpinMutex spin_lock;
|
|
|
|
function<void(size_t)> fun = [&](size_t i) {
|
|
spin_lock.lock();
|
|
count++;
|
|
size_t process = (size_t)ceil(count / s);
|
|
if (process > pc) {
|
|
pc = process;
|
|
stringstream msg;
|
|
msg << "\r" << setw(3) << setfill(' ') << process << "% " << left
|
|
<< setfill('=') << setw(process / 2 + 1) << "[" << right
|
|
<< setfill(' ') << setw(51 - process / 2) << "]";
|
|
cout << msg.str() << flush;
|
|
}
|
|
spin_lock.unlock();
|
|
|
|
auto query = linear_queries_[i];
|
|
|
|
FilterResultCache filter_cache;
|
|
std::shared_ptr<IndexFilter> filter_ptr = nullptr;
|
|
if (filter_mode_ == FM_TAG) {
|
|
if (batch_taglists_[i].size() != 1) {
|
|
LOG_ERROR("query tag list not equal to one!");
|
|
return;
|
|
}
|
|
|
|
int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[i][0],
|
|
tag_key_list_);
|
|
if (ret != 0) {
|
|
LOG_ERROR("prefilter failed, idx: %zu", i);
|
|
return;
|
|
}
|
|
|
|
auto filterFunc = [&](uint64_t key) {
|
|
return filter_cache.find(key);
|
|
};
|
|
|
|
filter_ptr = std::make_shared<IndexFilter>();
|
|
filter_ptr->set(filterFunc);
|
|
}
|
|
|
|
core_interface::DenseVector dense_query;
|
|
dense_query.data = query.data();
|
|
core_interface::VectorData query_data;
|
|
query_data.vector = dense_query;
|
|
|
|
auto query_param = std::make_shared<core_interface::FlatQueryParam>();
|
|
query_param->topk = gt_count;
|
|
query_param->is_linear = true;
|
|
query_param->filter = filter_ptr;
|
|
|
|
core_interface::SearchResult search_result;
|
|
int ret = index->Search(query_data, query_param, &search_result);
|
|
if (ret < 0) {
|
|
LOG_ERROR("Failed to linear search, ret=%d %s", ret,
|
|
IndexError::What(ret));
|
|
error.exchange(true);
|
|
return;
|
|
}
|
|
auto &result = search_result.doc_list_;
|
|
vector<pair<uint64_t, float>> one_gt;
|
|
one_gt.reserve(gt_count);
|
|
|
|
for (auto knn : result) {
|
|
one_gt.emplace_back(knn.key(), knn.score());
|
|
}
|
|
gt_[i] = one_gt;
|
|
};
|
|
for (size_t i = 0; i < linear_queries_.size(); ++i) {
|
|
if (error) {
|
|
break;
|
|
}
|
|
pool_->enqueue_and_wake(Closure::New(fun, i));
|
|
}
|
|
pool_->wait_finish();
|
|
|
|
if (error) {
|
|
cout << endl
|
|
<< "Ground truth file[" << gt_file << "] create failed!" << endl;
|
|
gtf.close();
|
|
remove(gt_file.c_str());
|
|
return false;
|
|
}
|
|
|
|
for (size_t i = 0; i < gt_.size(); ++i) {
|
|
auto > = gt_[i];
|
|
|
|
gtf.write(>_count, sizeof(int));
|
|
|
|
for (size_t j = 0; j < gt.size(); j++) {
|
|
auto &one_gt = gt[j];
|
|
|
|
gtf.write(&one_gt.first, sizeof(uint64_t));
|
|
gtf.write(&one_gt.second, sizeof(float));
|
|
}
|
|
|
|
// if ground truth is less than gt count, fill it up
|
|
if (gt.size() != gt_count) {
|
|
std::cout
|
|
<< "WARN: GT result count less than GT expected count, index: "
|
|
<< i << ", expected GT count: " << gt_count
|
|
<< ", actual GT count: " << gt.size() << std::endl;
|
|
|
|
uint64_t key{-1LLU};
|
|
float score{std::nanf("")};
|
|
|
|
for (size_t j = gt.size(); j < gt_count; ++j) {
|
|
gtf.write(&key, sizeof(uint64_t));
|
|
gtf.write(&score, sizeof(float));
|
|
}
|
|
}
|
|
}
|
|
|
|
gtf.close();
|
|
|
|
if (!File::Rename(gt_file_temp, gt_file)) {
|
|
LOG_ERROR("failed to rename ground truth file, src: %s, dst: %s",
|
|
gt_file_temp.c_str(), gt_file.c_str());
|
|
|
|
return false;
|
|
}
|
|
|
|
cout << endl
|
|
<< "Ground truth file create successful in "
|
|
<< timer.milli_seconds() / 1000 << "s." << endl;
|
|
} else {
|
|
if (!gtf.open(gt_file.c_str(), true)) {
|
|
LOG_ERROR("Failed to open ground truth file[%s]", gt_file.c_str());
|
|
return false;
|
|
}
|
|
size_t file_size = gtf.size();
|
|
|
|
constexpr size_t LENGTH = 10240;
|
|
constexpr size_t GT_PAIR_SIZE = sizeof(uint64_t) + sizeof(float);
|
|
|
|
char *buffer = new char[LENGTH];
|
|
gtf.read(buffer, sizeof(int));
|
|
|
|
size_t gt_count_input = (size_t) * (int *)buffer;
|
|
size_t one_query_line_size = sizeof(int) + GT_PAIR_SIZE * gt_count_input;
|
|
|
|
if (gt_count != gt_count_input || file_size % one_query_line_size != 0) {
|
|
LOG_ERROR("Ground truth file[%s] content error!", gt_file.c_str());
|
|
gtf.close();
|
|
return false;
|
|
}
|
|
|
|
size_t query_num = file_size / one_query_line_size;
|
|
if (one_query_line_size > LENGTH) {
|
|
delete[] buffer;
|
|
buffer = new char[one_query_line_size];
|
|
}
|
|
|
|
for (size_t n = 0; n < query_num; ++n) {
|
|
gtf.read(n * one_query_line_size, buffer, one_query_line_size);
|
|
vector<pair<uint64_t, float>> one_gt;
|
|
one_gt.reserve(gt_count);
|
|
|
|
for (size_t i = 0; i < gt_count; ++i) {
|
|
uint64_t key = *(uint64_t *)(buffer + sizeof(int) + GT_PAIR_SIZE * i);
|
|
float score = *(float *)(buffer + sizeof(int) + GT_PAIR_SIZE * i +
|
|
sizeof(uint64_t));
|
|
|
|
if (key != -1LLU) {
|
|
one_gt.emplace_back(key, score);
|
|
}
|
|
}
|
|
gt_.emplace_back(one_gt);
|
|
}
|
|
delete[] buffer;
|
|
cout << "Load ground truth file[" << gt_file << "] done!" << endl;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
|
|
void recall_one_dense(
|
|
core_interface::Index::Pointer index,
|
|
core_interface::BaseIndexQueryParam::Pointer query_param, size_t topk,
|
|
size_t idx,
|
|
std::vector<pair<std::fstream *, std::fstream *>> &output_fs) {
|
|
const auto &query = batch_queries_[idx];
|
|
|
|
size_t thread_index = pool_->indexof_this();
|
|
fstream *knn_fs = nullptr;
|
|
fstream *linear_fs = nullptr;
|
|
if (output_fs.size() > thread_index) {
|
|
knn_fs = output_fs[thread_index].first;
|
|
linear_fs = output_fs[thread_index].second;
|
|
}
|
|
|
|
auto cal_recall = [&, this](const std::vector<IndexDocument> &knn_res,
|
|
size_t query_idx) {
|
|
vector<IndexDocument> linear_res;
|
|
|
|
size_t result_size = std::min(topk, gt_[query_idx].size());
|
|
if (result_size == 0) {
|
|
return;
|
|
}
|
|
|
|
for (size_t i = 0; i < result_size; ++i) {
|
|
auto gt_node = gt_[query_idx][i];
|
|
|
|
linear_res.emplace_back(gt_node.first, gt_node.second, gt_node.first);
|
|
}
|
|
|
|
|
|
if (knn_fs) {
|
|
for (auto knn : knn_res) {
|
|
string str = "query[" + to_string(query_idx) + "]\tkey[" +
|
|
to_string(knn.key()) + "], dist[" +
|
|
to_string(knn.score()) + "]\n";
|
|
knn_fs->write(str.c_str(), str.size());
|
|
}
|
|
}
|
|
size_t match = 0;
|
|
bool asc =
|
|
(linear_res.size() > 1 &&
|
|
(linear_res[0].score() > linear_res[linear_res.size() - 1].score()))
|
|
? false
|
|
: true;
|
|
|
|
map<int32_t, size_t> topk_matchs;
|
|
if (g_compare_by_id) {
|
|
for (size_t i = 0; i < topk_ids_.size(); ++i) {
|
|
topk_matchs[topk_ids_[i]] = 0;
|
|
}
|
|
}
|
|
for (size_t i = 0, j = 0; i < linear_res.size();) {
|
|
bool m = false; // if current doc matched in max topk
|
|
bool changed = true; // if i changed
|
|
if (g_compare_by_id) {
|
|
for (size_t k = 0; k < topk_ids_.size(); ++k) {
|
|
size_t dynamic_size = (size_t)topk_ids_[k];
|
|
for (; dynamic_size + 1 < knn_res.size(); ++dynamic_size) {
|
|
if (fabs(knn_res[dynamic_size - 1].score() -
|
|
knn_res[dynamic_size].score()) >=
|
|
numeric_limits<float>::epsilon()) {
|
|
break;
|
|
}
|
|
}
|
|
for (size_t l = 0; l < dynamic_size && l < knn_res.size(); ++l) {
|
|
if (linear_res[i].key() == knn_res[l].key()) {
|
|
topk_matchs[topk_ids_[k]]++;
|
|
if (k == topk_ids_.size() - 1) {
|
|
m = true;
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
++i;
|
|
auto it = recall_res_.find(i);
|
|
if (it != recall_res_.end()) {
|
|
lock_guard<mutex> lock(recall_lock);
|
|
it->second += 100.0 * topk_matchs[i] / i;
|
|
}
|
|
} else {
|
|
size_t cur_topk = i + 1;
|
|
if (j < knn_res.size()) {
|
|
if (fabs(linear_res[i].score() - knn_res[j].score()) <
|
|
g_recall_precision) {
|
|
++j;
|
|
++i;
|
|
match++;
|
|
m = true;
|
|
} else {
|
|
if ((asc && linear_res[i].score() < knn_res[j].score()) ||
|
|
(!asc && linear_res[i].score() > knn_res[j].score())) {
|
|
++i;
|
|
} else {
|
|
changed = false;
|
|
++j;
|
|
}
|
|
}
|
|
} else {
|
|
++i;
|
|
}
|
|
auto it = recall_res_.find(cur_topk);
|
|
if (changed && it != recall_res_.end()) {
|
|
lock_guard<mutex> lock(recall_lock);
|
|
it->second += 100.0 * match / cur_topk;
|
|
}
|
|
}
|
|
if (linear_fs && changed) {
|
|
string str = string(m ? " HIT" : "NOT HIT") + " query[" +
|
|
to_string(idx) + "]\tkey[" +
|
|
to_string(linear_res[i - 1].key()) + "], dist[" +
|
|
to_string(linear_res[i - 1].score()) + "]\n";
|
|
linear_fs->write(str.c_str(), str.size());
|
|
}
|
|
}
|
|
};
|
|
|
|
// prefilter
|
|
FilterResultCache filter_cache;
|
|
std::shared_ptr<IndexFilter> filter_ptr = nullptr;
|
|
if (filter_mode_ == FM_TAG) {
|
|
if (batch_taglists_[idx].size() != 1) {
|
|
LOG_ERROR("query tag list not equal to one!");
|
|
return;
|
|
}
|
|
|
|
int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[idx][0],
|
|
tag_key_list_);
|
|
if (ret != 0) {
|
|
LOG_ERROR("prefilter failed, idx: %zu", idx);
|
|
return;
|
|
}
|
|
|
|
auto filterFunc = [&](uint64_t key) { return filter_cache.find(key); };
|
|
|
|
filter_ptr = std::make_shared<core::IndexFilter>();
|
|
filter_ptr->set(filterFunc);
|
|
}
|
|
|
|
core_interface::DenseVector dense_query;
|
|
dense_query.data = query.data();
|
|
core_interface::VectorData query_data;
|
|
query_data.vector = dense_query;
|
|
|
|
// query_param is required in the config, so it should not be nullptr
|
|
auto query_param_clone = query_param->Clone();
|
|
query_param_clone->topk = topk;
|
|
query_param_clone->filter = filter_ptr;
|
|
query_param_clone->is_linear = false;
|
|
|
|
if (call_batch_api_) {
|
|
size_t qnum = query.size() / dim_;
|
|
// For batch search, we need to search each query separately
|
|
// since Index::Search doesn't support batch natively in the same way
|
|
for (size_t i = 0; i < qnum; ++i) {
|
|
size_t query_idx = idx * batch_count_ + i;
|
|
if (query_idx >= linear_queries_.size()) {
|
|
break;
|
|
}
|
|
|
|
const auto &single_query = linear_queries_[query_idx];
|
|
core_interface::DenseVector single_dense_query;
|
|
single_dense_query.data = single_query.data();
|
|
core_interface::VectorData single_query_data;
|
|
single_query_data.vector = single_dense_query;
|
|
|
|
core_interface::SearchResult search_result;
|
|
int ret =
|
|
index->Search(single_query_data, query_param_clone, &search_result);
|
|
if (ret < 0) {
|
|
LOG_ERROR("Failed to knn_search batch, ret=%d %s", ret,
|
|
IndexError::What(ret));
|
|
return;
|
|
}
|
|
auto &knn_res = search_result.doc_list_;
|
|
cal_recall(knn_res, query_idx);
|
|
}
|
|
} else {
|
|
core_interface::SearchResult search_result;
|
|
int ret = index->Search(query_data, query_param_clone, &search_result);
|
|
if (ret < 0) {
|
|
LOG_ERROR("Failed to knn_search, ret=%d %s", ret,
|
|
IndexError::What(ret));
|
|
return;
|
|
}
|
|
auto &knn_res = search_result.doc_list_;
|
|
cal_recall(knn_res, idx);
|
|
}
|
|
|
|
// std::cout << "id: " << index << ": \n" <<
|
|
// knn_context->flow_context()->searcher_context()->profiler().display();
|
|
}
|
|
|
|
private:
|
|
IndexQueryMeta qmeta_{};
|
|
size_t threads_;
|
|
bool call_batch_api_;
|
|
string output_;
|
|
size_t batch_count_;
|
|
shared_ptr<ThreadPool> pool_;
|
|
|
|
// for gt
|
|
vector<vector<T>> linear_queries_;
|
|
vector<SparseData<T>> linear_sparse_data_;
|
|
vector<vector<uint64_t>> linear_taglists_;
|
|
|
|
// for recall
|
|
vector<vector<T>> batch_queries_;
|
|
vector<vector<uint32_t>> batch_sparse_counts_;
|
|
vector<vector<uint32_t>> batch_sparse_indices_;
|
|
vector<vector<T>> batch_sparse_features_;
|
|
vector<vector<vector<uint64_t>>> batch_taglists_;
|
|
|
|
size_t dim_;
|
|
size_t total_querys_;
|
|
|
|
map<size_t, float> recall_res_;
|
|
vector<int32_t> topk_ids_;
|
|
vector<vector<pair<uint64_t, float>>> gt_;
|
|
|
|
bool external_gt_file_enabled_{false};
|
|
|
|
FilterMode filter_mode_{FM_NONE};
|
|
|
|
static bool STOP_NOW;
|
|
|
|
// Tag lists for filtering
|
|
std::vector<std::vector<uint64_t>> id_to_tags_list_;
|
|
std::vector<uint64_t> tag_key_list_;
|
|
|
|
public:
|
|
void set_tag_lists(const std::vector<std::vector<uint64_t>> &id_to_tags_list,
|
|
const std::vector<uint64_t> &tag_key_list) {
|
|
id_to_tags_list_ = id_to_tags_list;
|
|
tag_key_list_ = tag_key_list;
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
bool Recall<T>::STOP_NOW = false;
|
|
|
|
//--------------------------------------------------
|
|
// Sparse Recall
|
|
//--------------------------------------------------
|
|
template <typename T>
|
|
class SparseRecall {
|
|
public:
|
|
SparseRecall(size_t threads, const string &output, size_t batch_count,
|
|
FilterMode filter_mode)
|
|
: threads_(threads),
|
|
output_(output),
|
|
batch_count_(batch_count),
|
|
filter_mode_{filter_mode} {
|
|
if (threads_ == 0) {
|
|
pool_ = make_shared<ThreadPool>(true);
|
|
threads_ = pool_->count();
|
|
cout << "Using cpu count as thread pool count[" << threads_ << "]"
|
|
<< endl;
|
|
} else {
|
|
pool_ = make_shared<ThreadPool>(threads_, true);
|
|
cout << "Using thread pool count[" << threads_ << "]" << endl;
|
|
}
|
|
if (batch_count_ < 1) {
|
|
batch_count_ = 1;
|
|
call_batch_api_ = false;
|
|
} else {
|
|
call_batch_api_ = true;
|
|
}
|
|
}
|
|
|
|
static void stop(int signo) {
|
|
if (STOP_NOW) {
|
|
exit(signo);
|
|
}
|
|
STOP_NOW = true;
|
|
cout << "\rTrying to stop. press [Ctrl+C] again kill immediately." << endl
|
|
<< flush;
|
|
}
|
|
|
|
int transform_queries_without_hybrid_scale(
|
|
const vector<vector<T>> &queries,
|
|
const vector<vector<T>> &sparse_features,
|
|
vector<vector<T>> *queries_output,
|
|
vector<vector<T>> *sparse_features_output) {
|
|
if (!queries_output || !sparse_features_output) {
|
|
LOG_ERROR("input should not be empty in transfrom queries");
|
|
|
|
return -1;
|
|
}
|
|
|
|
queries_output->clear();
|
|
sparse_features_output->clear();
|
|
|
|
for (size_t i = 0; i < queries.size(); ++i) {
|
|
vector<T> query_output;
|
|
vector<T> sparse_feature_output;
|
|
|
|
transform_query_without_hybrid_scale(queries[i], sparse_features[i],
|
|
&query_output,
|
|
&sparse_feature_output);
|
|
|
|
queries_output->push_back(query_output);
|
|
sparse_features_output->push_back(sparse_feature_output);
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
void run_sparse(core_interface::Index::Pointer index,
|
|
core_interface::BaseIndexQueryParam::Pointer query_param,
|
|
const string &recall_tops, size_t gt_count) {
|
|
StringHelper::Split(recall_tops, ",", &topk_ids_);
|
|
std::sort(topk_ids_.begin(), topk_ids_.end());
|
|
|
|
for (auto i : topk_ids_) {
|
|
recall_res_[i] = 0.0f;
|
|
}
|
|
size_t topk = recall_res_.rbegin()->first;
|
|
|
|
gt_count = topk < gt_count ? gt_count : topk;
|
|
|
|
if (external_gt_file_enabled_) {
|
|
cout << "Internal ground truth file NOT used since external ground truth "
|
|
"file has been loaded"
|
|
<< endl;
|
|
} else {
|
|
cout << "Loading internal ground truth file" << endl;
|
|
|
|
if (!load_gt_sparse(index, gt_count)) {
|
|
LOG_ERROR("Load ground truth file failed!");
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (batch_sparse_counts_.size() < threads_) {
|
|
threads_ = batch_sparse_counts_.size();
|
|
pool_ = make_shared<ThreadPool>(true, threads_);
|
|
cout << "Query size too small, resize thread pool count[" << threads_
|
|
<< "]" << endl;
|
|
}
|
|
|
|
// Prepare file handler
|
|
vector<pair<fstream *, fstream *>> output_fs;
|
|
if (!output_.empty()) {
|
|
string cmd = "mkdir -p " + output_;
|
|
int ret = system(cmd.c_str());
|
|
if (ret != 0) {
|
|
LOG_ERROR("execute cmd %s failed, ret=%d", cmd.c_str(), ret);
|
|
return;
|
|
}
|
|
struct stat sb;
|
|
if (stat(output_.c_str(), &sb) == 0 && S_ISDIR(sb.st_mode)) {
|
|
cout << "logs output to : " << output_ << endl;
|
|
for (size_t i = 0; i < threads_; ++i) {
|
|
fstream *fs_k = new fstream();
|
|
fs_k->open(output_ + "/t" + to_string(i) + ".knn", ios::out);
|
|
fstream *fs_l = new fstream();
|
|
fs_l->open(output_ + "/t" + to_string(i) + ".linear", ios::out);
|
|
output_fs.push_back(make_pair(fs_k, fs_l));
|
|
}
|
|
}
|
|
}
|
|
|
|
signal(SIGINT, stop);
|
|
size_t i = 0;
|
|
for (; !STOP_NOW && i < batch_sparse_counts_.size();) {
|
|
if (pool_->pending_count() >= pool_->count()) {
|
|
this_thread::sleep_for(chrono::microseconds(1));
|
|
continue;
|
|
}
|
|
|
|
Closure::Pointer task =
|
|
Closure::New(this, &SparseRecall::recall_one_sparse, index,
|
|
query_param, topk, i, output_fs);
|
|
pool_->enqueue_and_wake(task);
|
|
|
|
i++;
|
|
}
|
|
pool_->wait_finish();
|
|
|
|
for (auto fs : output_fs) {
|
|
fs.first->close();
|
|
fs.second->close();
|
|
delete fs.first;
|
|
delete fs.second;
|
|
}
|
|
cout << "Process query: " << i << endl;
|
|
for (auto it : recall_res_) {
|
|
cout << "Recall@" << it.first << ": "
|
|
<< it.second / linear_queries_.size() << endl;
|
|
}
|
|
}
|
|
|
|
bool load_query(const std::string &query_file, const std::string &first_sep,
|
|
const std::string &second_sep) {
|
|
TxtInputReader<T> reader;
|
|
|
|
if (!reader.load_query(query_file, first_sep, second_sep, linear_queries_,
|
|
linear_sparse_data_, linear_taglists_)) {
|
|
LOG_ERROR("Load query error");
|
|
return false;
|
|
}
|
|
|
|
if (batch_count_ == 1) {
|
|
for (size_t i = 0; i < linear_sparse_data_.size(); ++i) {
|
|
vector<uint32_t> sparse_count;
|
|
sparse_count.push_back(linear_sparse_data_[i].count);
|
|
|
|
batch_sparse_counts_.push_back(sparse_count);
|
|
batch_sparse_indices_.push_back(linear_sparse_data_[i].indices);
|
|
batch_sparse_features_.push_back(linear_sparse_data_[i].features);
|
|
}
|
|
} else {
|
|
size_t num_batch =
|
|
(linear_queries_.size() + batch_count_ - 1) / batch_count_;
|
|
size_t idx = 0;
|
|
for (size_t n = 0; n < num_batch; ++n) {
|
|
vector<uint32_t> batch_sparse_count;
|
|
vector<uint32_t> batch_sparse_indices;
|
|
vector<T> batch_sparse_feature;
|
|
|
|
for (size_t i = 0; i < batch_count_; ++i) {
|
|
batch_sparse_count.push_back(linear_sparse_data_[idx].count);
|
|
|
|
for (size_t k = 0; k < linear_sparse_data_[idx].indices.size(); ++k) {
|
|
batch_sparse_indices.push_back(linear_sparse_data_[idx].indices[k]);
|
|
}
|
|
|
|
for (size_t k = 0; k < linear_sparse_data_[idx].features.size();
|
|
++k) {
|
|
batch_sparse_feature.push_back(
|
|
linear_sparse_data_[idx].features[k]);
|
|
}
|
|
|
|
idx = (idx + 1) % linear_queries_.size();
|
|
}
|
|
batch_sparse_counts_.push_back(batch_sparse_count);
|
|
batch_sparse_indices_.push_back(batch_sparse_indices);
|
|
batch_sparse_features_.push_back(batch_sparse_feature);
|
|
}
|
|
}
|
|
|
|
total_querys_ = linear_queries_.size();
|
|
if (typeid(T) == typeid(float)) {
|
|
qmeta_.set_data_type(IndexMeta::DataType::DT_FP32);
|
|
} else if (typeid(T) == typeid(int8_t)) {
|
|
qmeta_.set_data_type(IndexMeta::DataType::DT_INT8);
|
|
} else {
|
|
LOG_ERROR("unsupported type");
|
|
return false;
|
|
}
|
|
|
|
cout << "Load query done!" << endl;
|
|
return true;
|
|
}
|
|
|
|
bool load_gt_sparse(core_interface::Index::Pointer index, size_t gt_count) {
|
|
std::string crc_str = compute_crc();
|
|
|
|
string gt_file = string("gt.") + crc_str;
|
|
|
|
File gtf;
|
|
if (!gtf.IsRegular(gt_file.c_str())) {
|
|
cout << "Ground truth file[" << gt_file << "] not exist, try to create it"
|
|
<< endl;
|
|
ElapsedTime timer;
|
|
size_t size = sizeof(uint64_t) + sizeof(float);
|
|
size_t file_size =
|
|
linear_sparse_data_.size() * (sizeof(int) + size * gt_count);
|
|
|
|
std::string gt_file_temp = gt_file + ".tmp";
|
|
gtf.create(gt_file_temp.c_str(), file_size);
|
|
|
|
gt_.resize(linear_sparse_data_.size());
|
|
|
|
atomic_bool error(false);
|
|
size_t count = 0;
|
|
float s = linear_sparse_data_.size() / 100.0;
|
|
size_t pc = 0;
|
|
SpinMutex spin_lock;
|
|
|
|
function<void(size_t)> fun = [&](size_t i) {
|
|
spin_lock.lock();
|
|
count++;
|
|
size_t process = (size_t)ceil(count / s);
|
|
if (process > pc) {
|
|
pc = process;
|
|
stringstream msg;
|
|
msg << "\r" << setw(3) << setfill(' ') << process << "% " << left
|
|
<< setfill('=') << setw(process / 2 + 1) << "[" << right
|
|
<< setfill(' ') << setw(51 - process / 2) << "]";
|
|
cout << msg.str() << flush;
|
|
}
|
|
spin_lock.unlock();
|
|
|
|
SparseData<T> sparse_data = linear_sparse_data_[i];
|
|
|
|
// prefilter
|
|
FilterResultCache filter_cache;
|
|
std::shared_ptr<IndexFilter> filter_ptr = nullptr;
|
|
if (filter_mode_ == FM_TAG) {
|
|
if (batch_taglists_[i].size() != 1) {
|
|
LOG_ERROR("query tag list not equal to one!");
|
|
return;
|
|
}
|
|
|
|
int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[i][0],
|
|
tag_key_list_);
|
|
if (ret != 0) {
|
|
LOG_ERROR("prefilter failed, idx: %zu", i);
|
|
return;
|
|
}
|
|
|
|
auto filterFunc = [&](uint64_t key) {
|
|
return filter_cache.find(key);
|
|
};
|
|
|
|
filter_ptr = std::make_shared<IndexFilter>();
|
|
filter_ptr->set(filterFunc);
|
|
}
|
|
|
|
core_interface::SparseVector sparse_query;
|
|
sparse_query.count = sparse_data.count;
|
|
sparse_query.indices = sparse_data.indices.data();
|
|
sparse_query.values = sparse_data.features.data();
|
|
core_interface::VectorData query_data;
|
|
query_data.vector = sparse_query;
|
|
|
|
auto query_param = std::make_shared<core_interface::FlatQueryParam>();
|
|
query_param->topk = gt_count;
|
|
query_param->is_linear = true;
|
|
query_param->filter = filter_ptr;
|
|
|
|
core_interface::SearchResult search_result;
|
|
int ret = index->Search(query_data, query_param, &search_result);
|
|
if (ret < 0) {
|
|
LOG_ERROR("Failed to sparse linear search, ret=%d", ret);
|
|
error.exchange(true);
|
|
return;
|
|
}
|
|
auto &result = search_result.doc_list_;
|
|
|
|
vector<pair<uint64_t, float>> one_gt;
|
|
one_gt.reserve(gt_count);
|
|
|
|
for (auto knn : result) {
|
|
one_gt.emplace_back(knn.key(), knn.score());
|
|
}
|
|
gt_[i] = one_gt;
|
|
};
|
|
|
|
for (size_t i = 0; i < linear_sparse_data_.size(); ++i) {
|
|
if (error) {
|
|
break;
|
|
}
|
|
pool_->enqueue_and_wake(Closure::New(fun, i));
|
|
}
|
|
pool_->wait_finish();
|
|
|
|
if (error) {
|
|
cout << endl
|
|
<< "Ground truth file[" << gt_file << "] create failed!" << endl;
|
|
gtf.close();
|
|
remove(gt_file.c_str());
|
|
return false;
|
|
}
|
|
|
|
for (size_t i = 0; i < gt_.size(); ++i) {
|
|
auto > = gt_[i];
|
|
|
|
gtf.write(>_count, sizeof(int));
|
|
|
|
for (size_t j = 0; j < gt.size(); j++) {
|
|
auto &one_gt = gt[j];
|
|
|
|
gtf.write(&one_gt.first, sizeof(uint64_t));
|
|
gtf.write(&one_gt.second, sizeof(float));
|
|
}
|
|
|
|
// if ground truth is less than gt count, fill it up
|
|
if (gt.size() != gt_count) {
|
|
std::cout
|
|
<< "WARN: GT result count less than GT expected count, index: "
|
|
<< i << ", expected GT count: " << gt_count
|
|
<< ", actual GT count: " << gt.size() << std::endl;
|
|
|
|
uint64_t key{-1LLU};
|
|
float score{std::nanf("")};
|
|
|
|
for (size_t j = gt.size(); j < gt_count; ++j) {
|
|
gtf.write(&key, sizeof(uint64_t));
|
|
gtf.write(&score, sizeof(float));
|
|
}
|
|
}
|
|
}
|
|
gtf.close();
|
|
|
|
if (!File::Rename(gt_file_temp, gt_file)) {
|
|
LOG_ERROR("failed to rename ground truth file, src: %s, dst: %s",
|
|
gt_file_temp.c_str(), gt_file.c_str());
|
|
|
|
return false;
|
|
}
|
|
|
|
cout << endl
|
|
<< "Ground truth file create successful in "
|
|
<< timer.milli_seconds() / 1000 << "s." << endl;
|
|
} else {
|
|
if (!gtf.open(gt_file.c_str(), true)) {
|
|
LOG_ERROR("Failed to open ground truth file[%s]", gt_file.c_str());
|
|
return false;
|
|
}
|
|
size_t file_size = gtf.size();
|
|
|
|
constexpr size_t LENGTH = 10240;
|
|
constexpr size_t GT_PAIR_SIZE = sizeof(uint64_t) + sizeof(float);
|
|
|
|
char *buffer = new char[LENGTH];
|
|
gtf.read(buffer, sizeof(int));
|
|
|
|
size_t gt_count_input = (size_t) * (int *)buffer;
|
|
size_t one_query_line_size = sizeof(int) + GT_PAIR_SIZE * gt_count_input;
|
|
|
|
if (gt_count != gt_count_input || file_size % one_query_line_size != 0) {
|
|
LOG_ERROR("Ground truth file[%s] content error!", gt_file.c_str());
|
|
gtf.close();
|
|
return false;
|
|
}
|
|
|
|
size_t query_num = file_size / one_query_line_size;
|
|
if (one_query_line_size > LENGTH) {
|
|
delete[] buffer;
|
|
buffer = new char[one_query_line_size];
|
|
}
|
|
|
|
for (size_t n = 0; n < query_num; ++n) {
|
|
gtf.read(n * one_query_line_size, buffer, one_query_line_size);
|
|
vector<pair<uint64_t, float>> one_gt;
|
|
one_gt.reserve(gt_count);
|
|
|
|
for (size_t i = 0; i < gt_count; ++i) {
|
|
uint64_t key = *(uint64_t *)(buffer + sizeof(int) + GT_PAIR_SIZE * i);
|
|
float score = *(float *)(buffer + sizeof(int) + GT_PAIR_SIZE * i +
|
|
sizeof(uint64_t));
|
|
|
|
if (key != -1LLU) {
|
|
one_gt.emplace_back(key, score);
|
|
}
|
|
}
|
|
|
|
gt_.emplace_back(one_gt);
|
|
}
|
|
|
|
delete[] buffer;
|
|
cout << "Load ground truth file[" << gt_file << "] done!" << endl;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
bool load_external_gt_file(const std::string &external_gt_file,
|
|
const std::string &first_sep,
|
|
const std::string &second_sep) {
|
|
TxtInputReader<T> reader;
|
|
bool ret =
|
|
reader.load_external_gt(external_gt_file, first_sep, second_sep, gt_);
|
|
if (ret) {
|
|
cout << "Load external ground truth file["
|
|
<< File::BaseName(external_gt_file) << "] done!" << endl;
|
|
external_gt_file_enabled_ = true;
|
|
} else {
|
|
LOG_ERROR("Failed to load ground truth file!");
|
|
}
|
|
|
|
return ret;
|
|
}
|
|
|
|
private:
|
|
std::string compute_crc() {
|
|
uint32_t crc = 0u;
|
|
// sparse
|
|
if (linear_sparse_data_.size() > 0) {
|
|
for (size_t i = 0; i < linear_sparse_data_.size(); ++i) {
|
|
crc = Crc32c::Hash(&(linear_sparse_data_[i].count), sizeof(uint32_t),
|
|
crc);
|
|
crc =
|
|
Crc32c::Hash(linear_sparse_data_[i].indices.data(),
|
|
linear_sparse_data_[i].count * sizeof(uint32_t), crc);
|
|
crc = Crc32c::Hash(linear_sparse_data_[i].features.data(),
|
|
linear_sparse_data_[i].count * sizeof(T), crc);
|
|
}
|
|
}
|
|
|
|
char crc_str[64];
|
|
snprintf(crc_str, sizeof(crc_str), "%X", crc);
|
|
|
|
return std::string(crc_str);
|
|
}
|
|
|
|
|
|
void recall_one_sparse(
|
|
core_interface::Index::Pointer index,
|
|
core_interface::BaseIndexQueryParam::Pointer query_param, size_t topk,
|
|
size_t idx,
|
|
std::vector<pair<std::fstream *, std::fstream *>> &output_fs) {
|
|
const auto &sparse_count = batch_sparse_counts_[idx];
|
|
const auto &sparse_index = batch_sparse_indices_[idx];
|
|
const auto &sparse_feature = batch_sparse_features_[idx];
|
|
|
|
size_t thread_index = pool_->indexof_this();
|
|
fstream *knn_fs = nullptr;
|
|
fstream *linear_fs = nullptr;
|
|
if (output_fs.size() > thread_index) {
|
|
knn_fs = output_fs[thread_index].first;
|
|
linear_fs = output_fs[thread_index].second;
|
|
}
|
|
|
|
auto cal_recall = [&, this](const std::vector<IndexDocument> &knn_res,
|
|
size_t query_idx) {
|
|
vector<IndexDocument> linear_res;
|
|
|
|
size_t result_size = std::min(topk, gt_[query_idx].size());
|
|
if (result_size == 0) {
|
|
return;
|
|
}
|
|
|
|
for (size_t i = 0; i < result_size; ++i) {
|
|
auto gt_node = gt_[query_idx][i];
|
|
|
|
linear_res.emplace_back(gt_node.first, gt_node.second, gt_node.first);
|
|
}
|
|
|
|
if (knn_fs) {
|
|
for (auto knn : knn_res) {
|
|
string str = "query[" + to_string(query_idx) + "]\tkey[" +
|
|
to_string(knn.key()) + "], dist[" +
|
|
to_string(knn.score()) + "]\n";
|
|
knn_fs->write(str.c_str(), str.size());
|
|
}
|
|
}
|
|
|
|
size_t match = 0;
|
|
bool asc =
|
|
(linear_res.size() > 1 &&
|
|
(linear_res[0].score() > linear_res[linear_res.size() - 1].score()))
|
|
? false
|
|
: true;
|
|
|
|
map<int32_t, size_t> topk_matchs;
|
|
if (g_compare_by_id) {
|
|
for (size_t i = 0; i < topk_ids_.size(); ++i) {
|
|
topk_matchs[topk_ids_[i]] = 0;
|
|
}
|
|
}
|
|
|
|
for (size_t i = 0, j = 0; i < linear_res.size();) {
|
|
bool m = false; // if current doc matched in max topk
|
|
bool changed = true; // if i changed
|
|
if (g_compare_by_id) {
|
|
for (size_t k = 0; k < topk_ids_.size(); ++k) {
|
|
size_t dynamic_size = (size_t)topk_ids_[k];
|
|
for (; dynamic_size + 1 < knn_res.size(); ++dynamic_size) {
|
|
if (fabs(knn_res[dynamic_size - 1].score() -
|
|
knn_res[dynamic_size].score()) >=
|
|
numeric_limits<float>::epsilon()) {
|
|
break;
|
|
}
|
|
}
|
|
for (size_t l = 0; l < dynamic_size && l < knn_res.size(); ++l) {
|
|
if (linear_res[i].key() == knn_res[l].key()) {
|
|
topk_matchs[topk_ids_[k]]++;
|
|
if (k == topk_ids_.size() - 1) {
|
|
m = true;
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
++i;
|
|
|
|
auto it = recall_res_.find(i);
|
|
if (it != recall_res_.end()) {
|
|
lock_guard<mutex> lock(recall_lock);
|
|
it->second += 100.0 * topk_matchs[i] / i;
|
|
}
|
|
} else {
|
|
size_t cur_topk = i + 1;
|
|
if (j < knn_res.size()) {
|
|
if (fabs(linear_res[i].score() - knn_res[j].score()) <
|
|
g_recall_precision) {
|
|
++j;
|
|
++i;
|
|
match++;
|
|
m = true;
|
|
} else {
|
|
if ((asc && linear_res[i].score() < knn_res[j].score()) ||
|
|
(!asc && linear_res[i].score() > knn_res[j].score())) {
|
|
++i;
|
|
} else {
|
|
changed = false;
|
|
++j;
|
|
}
|
|
}
|
|
} else {
|
|
++i;
|
|
}
|
|
|
|
auto it = recall_res_.find(cur_topk);
|
|
if (changed && it != recall_res_.end()) {
|
|
lock_guard<mutex> lock(recall_lock);
|
|
it->second += 100.0 * match / cur_topk;
|
|
}
|
|
}
|
|
|
|
if (linear_fs && changed) {
|
|
string str = string(m ? " HIT" : "NOT HIT") + " query[" +
|
|
to_string(idx) + "]\tkey[" +
|
|
to_string(linear_res[i - 1].key()) + "], dist[" +
|
|
to_string(linear_res[i - 1].score()) + "]\n";
|
|
linear_fs->write(str.c_str(), str.size());
|
|
}
|
|
}
|
|
};
|
|
|
|
FilterResultCache filter_cache;
|
|
std::shared_ptr<IndexFilter> filter_ptr = nullptr;
|
|
if (filter_mode_ == FM_TAG) {
|
|
if (batch_taglists_[idx].size() != 1) {
|
|
LOG_ERROR("query tag list not equal to one!");
|
|
return;
|
|
}
|
|
|
|
int ret = filter_cache.filter(id_to_tags_list_, batch_taglists_[idx][0],
|
|
tag_key_list_);
|
|
if (ret != 0) {
|
|
LOG_ERROR("prefilter failed, idx: %zu", idx);
|
|
return;
|
|
}
|
|
|
|
auto filterFunc = [&](uint64_t key) { return filter_cache.find(key); };
|
|
|
|
filter_ptr = std::make_shared<core::IndexFilter>();
|
|
filter_ptr->set(filterFunc);
|
|
}
|
|
|
|
core_interface::SparseVector sparse_query;
|
|
sparse_query.count = sparse_count[0];
|
|
sparse_query.indices = sparse_index.data();
|
|
sparse_query.values = sparse_feature.data();
|
|
core_interface::VectorData query_data;
|
|
query_data.vector = sparse_query;
|
|
|
|
auto query_param_clone = query_param->Clone();
|
|
query_param_clone->topk = topk;
|
|
query_param_clone->filter = filter_ptr;
|
|
query_param_clone->is_linear = true;
|
|
|
|
if (call_batch_api_) {
|
|
// For batch search, we need to search each query separately
|
|
for (size_t i = 0; i < sparse_count.size(); ++i) {
|
|
size_t query_idx = idx * batch_count_ + i;
|
|
if (query_idx >= linear_sparse_data_.size()) {
|
|
break;
|
|
}
|
|
|
|
const auto &single_sparse = linear_sparse_data_[query_idx];
|
|
core_interface::SparseVector single_sparse_query;
|
|
single_sparse_query.count = single_sparse.count;
|
|
single_sparse_query.indices = single_sparse.indices.data();
|
|
single_sparse_query.values = single_sparse.features.data();
|
|
core_interface::VectorData single_query_data;
|
|
single_query_data.vector = single_sparse_query;
|
|
|
|
core_interface::SearchResult search_result;
|
|
int ret =
|
|
index->Search(single_query_data, query_param_clone, &search_result);
|
|
if (ret < 0) {
|
|
LOG_ERROR("Failed to sparse_knn_search batch, ret=%d %s", ret,
|
|
IndexError::What(ret));
|
|
return;
|
|
}
|
|
auto &knn_res = search_result.doc_list_;
|
|
cal_recall(knn_res, query_idx);
|
|
}
|
|
} else {
|
|
core_interface::SearchResult search_result;
|
|
int ret = index->Search(query_data, query_param_clone, &search_result);
|
|
if (ret < 0) {
|
|
LOG_ERROR("Failed to sparse_knn_search, ret=%d %s", ret,
|
|
IndexError::What(ret));
|
|
return;
|
|
}
|
|
auto &knn_res = search_result.doc_list_;
|
|
cal_recall(knn_res, idx);
|
|
}
|
|
}
|
|
|
|
private:
|
|
IndexQueryMeta qmeta_{};
|
|
size_t threads_;
|
|
bool call_batch_api_;
|
|
string output_;
|
|
size_t batch_count_;
|
|
shared_ptr<ThreadPool> pool_;
|
|
|
|
// for gt
|
|
vector<vector<T>> linear_queries_;
|
|
vector<SparseData<T>> linear_sparse_data_;
|
|
vector<uint32_t> linear_partitions_;
|
|
vector<vector<uint64_t>> linear_taglists_;
|
|
|
|
std::map<std::string, vector<vector<T>>> linear_queries_scaled_;
|
|
std::map<std::string, vector<vector<T>>> linear_sparse_features_scaled_;
|
|
|
|
// for recall
|
|
vector<vector<T>> batch_queries_;
|
|
vector<vector<uint32_t>> batch_sparse_counts_;
|
|
vector<vector<uint32_t>> batch_sparse_indices_;
|
|
vector<vector<T>> batch_sparse_features_;
|
|
vector<vector<uint32_t>> batch_partitions_;
|
|
vector<vector<vector<uint64_t>>> batch_taglists_;
|
|
|
|
std::map<std::string, vector<vector<T>>> batch_queries_scaled_;
|
|
std::map<std::string, vector<vector<T>>> batch_sparse_features_scaled_;
|
|
|
|
size_t total_querys_;
|
|
|
|
map<size_t, float> recall_res_;
|
|
vector<int32_t> topk_ids_;
|
|
vector<vector<pair<uint64_t, float>>> gt_;
|
|
|
|
map<string, vector<vector<pair<uint64_t, float>>>> gt_hybrid_;
|
|
bool external_gt_file_enabled_{false};
|
|
|
|
FilterMode filter_mode_{FM_NONE};
|
|
|
|
// Tag lists for filtering
|
|
std::vector<std::vector<uint64_t>> id_to_tags_list_;
|
|
std::vector<uint64_t> tag_key_list_;
|
|
|
|
public:
|
|
void set_tag_lists(const std::vector<std::vector<uint64_t>> &id_to_tags_list,
|
|
const std::vector<uint64_t> &tag_key_list) {
|
|
id_to_tags_list_ = id_to_tags_list;
|
|
tag_key_list_ = tag_key_list;
|
|
}
|
|
|
|
static bool STOP_NOW;
|
|
};
|
|
|
|
template <typename T>
|
|
bool SparseRecall<T>::STOP_NOW = false;
|
|
|
|
bool check_config(YAML::Node &config_node) {
|
|
auto common = config_node["IndexCommon"];
|
|
if (!common) {
|
|
LOG_ERROR("Can not find [IndexCommon] in config");
|
|
return false;
|
|
}
|
|
if (!common["IndexConfig"]) {
|
|
LOG_ERROR("Can not find [IndexConfig] in config");
|
|
return false;
|
|
}
|
|
if (!common["IndexPath"]) {
|
|
LOG_ERROR("Can not find [IndexPath] in config");
|
|
return false;
|
|
}
|
|
if (!common["TopK"]) {
|
|
LOG_ERROR("Can not find [TopK] in config");
|
|
return false;
|
|
}
|
|
if (!common["QueryFile"]) {
|
|
LOG_ERROR("Can not find [QueryFile] in config");
|
|
return false;
|
|
}
|
|
|
|
auto query_config = config_node["QueryConfig"];
|
|
if (!query_config) {
|
|
LOG_ERROR("Can not find [QueryConfig] in config");
|
|
return false;
|
|
}
|
|
if (!query_config["QueryParam"]) {
|
|
LOG_ERROR("Can not find [QueryConfig.QueryParam] in config");
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
void usage(void) {
|
|
cout << "Usage: recall CONFIG.yaml [plugin file path]" << endl;
|
|
}
|
|
|
|
int recall_dense(std::string &query_type, size_t thread_count,
|
|
size_t batch_count, string top_k, size_t gt_count,
|
|
string query_file, string &first_sep, string &second_sep,
|
|
string &ground_truth_file, string &ground_truth_first_sep,
|
|
string ground_truth_second_sep,
|
|
core_interface::Index::Pointer index,
|
|
core_interface::BaseIndexQueryParam::Pointer query_param,
|
|
string &index_dir, string &log_dir, FilterMode filter_mode) {
|
|
std::vector<std::vector<uint64_t>> id_to_tags_list;
|
|
std::vector<uint64_t> tag_key_list;
|
|
// Load tag lists if available
|
|
load_taglists(index_dir, id_to_tags_list, tag_key_list);
|
|
|
|
if (query_type == "float") {
|
|
Recall<float> recall(thread_count, log_dir, batch_count, filter_mode);
|
|
if (!recall.load_query(query_file, first_sep, second_sep)) {
|
|
return -1;
|
|
}
|
|
|
|
recall.set_tag_lists(id_to_tags_list, tag_key_list);
|
|
|
|
if (ground_truth_file != "") {
|
|
if (!recall.load_external_gt_file(ground_truth_file,
|
|
ground_truth_first_sep,
|
|
ground_truth_second_sep)) {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
recall.run_dense(index, query_param, top_k, gt_count);
|
|
} else if (query_type == "int8") {
|
|
Recall<int8_t> recall(thread_count, log_dir, batch_count, filter_mode);
|
|
if (!recall.load_query(query_file, first_sep, second_sep)) {
|
|
return -1;
|
|
}
|
|
|
|
recall.set_tag_lists(id_to_tags_list, tag_key_list);
|
|
|
|
if (ground_truth_file != "") {
|
|
if (!recall.load_external_gt_file(ground_truth_file,
|
|
ground_truth_first_sep,
|
|
ground_truth_second_sep)) {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
recall.run_dense(index, query_param, top_k, gt_count);
|
|
} else if (query_type == "binary") {
|
|
Recall<uint32_t> recall(thread_count, log_dir, batch_count, filter_mode);
|
|
if (!recall.load_query(query_file, first_sep, second_sep)) {
|
|
return -1;
|
|
}
|
|
|
|
recall.set_tag_lists(id_to_tags_list, tag_key_list);
|
|
|
|
if (ground_truth_file != "") {
|
|
if (!recall.load_external_gt_file(ground_truth_file,
|
|
ground_truth_first_sep,
|
|
ground_truth_second_sep)) {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
recall.run_dense(index, query_param, top_k, gt_count);
|
|
} else if (query_type == "binary64") {
|
|
Recall<uint64_t> recall(thread_count, log_dir, batch_count, filter_mode);
|
|
if (!recall.load_query(query_file, first_sep, second_sep)) {
|
|
return -1;
|
|
}
|
|
|
|
recall.set_tag_lists(id_to_tags_list, tag_key_list);
|
|
|
|
if (ground_truth_file != "") {
|
|
if (!recall.load_external_gt_file(ground_truth_file,
|
|
ground_truth_first_sep,
|
|
ground_truth_second_sep)) {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
recall.run_dense(index, query_param, top_k, gt_count);
|
|
} else {
|
|
LOG_ERROR("Can not recognize type: %s", query_type.c_str());
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
int recall_sparse(std::string &query_type, size_t thread_count,
|
|
size_t batch_count, string top_k, size_t gt_count,
|
|
string &query_file, string &first_sep, string &second_sep,
|
|
string &ground_truth_file, string &ground_truth_first_sep,
|
|
string &ground_truth_second_sep,
|
|
core_interface::Index::Pointer index,
|
|
core_interface::BaseIndexQueryParam::Pointer query_param,
|
|
string &index_dir, string &log_dir, FilterMode filter_mode) {
|
|
if (query_type == "float") {
|
|
SparseRecall<float> recall(thread_count, log_dir, batch_count, filter_mode);
|
|
if (!recall.load_query(query_file, first_sep, second_sep)) {
|
|
return -1;
|
|
}
|
|
|
|
if (ground_truth_file != "") {
|
|
if (!recall.load_external_gt_file(ground_truth_file,
|
|
ground_truth_first_sep,
|
|
ground_truth_second_sep)) {
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
std::vector<std::vector<uint64_t>> id_to_tags_list;
|
|
std::vector<uint64_t> tag_key_list;
|
|
// Load tag lists if available
|
|
if (load_taglists(index_dir, id_to_tags_list, tag_key_list) != 0) {
|
|
LOG_ERROR("Failed to load tag lists");
|
|
return -1;
|
|
}
|
|
|
|
recall.set_tag_lists(id_to_tags_list, tag_key_list);
|
|
|
|
recall.run_sparse(index, query_param, top_k, gt_count);
|
|
} else {
|
|
LOG_ERROR("Can not recognize type: %s", query_type.c_str());
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
int get_recall_precision(string &recall_precision_string) {
|
|
constexpr float DEFAULT_RECALL_PRECISION = 1e-6;
|
|
|
|
if (recall_precision_string == "") {
|
|
g_recall_precision = DEFAULT_RECALL_PRECISION;
|
|
return true;
|
|
}
|
|
|
|
try {
|
|
g_recall_precision = std::stof(recall_precision_string);
|
|
std::cout << "Recall Score Precesion: " << g_recall_precision << std::endl;
|
|
} catch (const std::invalid_argument &e) {
|
|
LOG_ERROR("Exeception in getting recall precision: %s, value: %s", e.what(),
|
|
recall_precision_string.c_str());
|
|
return -1;
|
|
} catch (const std::out_of_range &e) {
|
|
LOG_ERROR(
|
|
"Out of range exception in getting recall precision: %s, value: %s",
|
|
e.what(), recall_precision_string.c_str());
|
|
return -1;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
int main(int argc, char *argv[]) {
|
|
if (argc < 2) {
|
|
usage();
|
|
return -1;
|
|
}
|
|
|
|
IndexPluginBroker broker;
|
|
std::string error;
|
|
for (int i = 2; i < argc; ++i) {
|
|
if (!broker.emplace(argv[i], &error)) {
|
|
LOG_ERROR("Failed to load plugin: %s (%s)", argv[i], error.c_str());
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
YAML::Node config_node;
|
|
try {
|
|
config_node = YAML::LoadFile(argv[1]);
|
|
} catch (...) {
|
|
LOG_ERROR("Load YAML file[%s] failed!", argv[1]);
|
|
return -1;
|
|
}
|
|
if (!check_config(config_node)) {
|
|
return -1;
|
|
}
|
|
auto config_common = config_node["IndexCommon"];
|
|
|
|
map<string, int> LOG_LEVEL = {{"debug", IndexLogger::LEVEL_DEBUG},
|
|
{"info", IndexLogger::LEVEL_INFO},
|
|
{"warn", IndexLogger::LEVEL_WARN},
|
|
{"error", IndexLogger::LEVEL_ERROR},
|
|
{"fatal", IndexLogger::LEVEL_FATAL}};
|
|
string log_level = config_common["LogLevel"]
|
|
? config_common["LogLevel"].as<string>()
|
|
: "debug";
|
|
transform(log_level.begin(), log_level.end(), log_level.begin(), ::tolower);
|
|
if (LOG_LEVEL.find(log_level) != LOG_LEVEL.end()) {
|
|
IndexLoggerBroker::SetLevel(LOG_LEVEL[log_level]);
|
|
zvec::ailego::LoggerBroker::SetLevel(LOG_LEVEL[log_level]);
|
|
}
|
|
|
|
// Calculate Recall
|
|
string log_dir = "";
|
|
if (config_common["RecallLogDir"]) {
|
|
log_dir = config_common["RecallLogDir"].as<string>();
|
|
}
|
|
size_t thread_count = config_common["RecallThreadCount"]
|
|
? config_common["RecallThreadCount"].as<uint64_t>()
|
|
: 0;
|
|
size_t gt_count = config_common["RecallGTCount"]
|
|
? config_common["RecallGTCount"].as<uint64_t>()
|
|
: 100;
|
|
size_t batch_count = config_common["RecallBatchCount"]
|
|
? config_common["RecallBatchCount"].as<uint64_t>()
|
|
: 0;
|
|
g_compare_by_id = config_common["CompareById"]
|
|
? config_common["CompareById"].as<bool>()
|
|
: 0;
|
|
string top_k = config_common["TopK"].as<string>();
|
|
|
|
string recall_precision_string =
|
|
config_common["RecallScorePrecision"]
|
|
? config_common["RecallScorePrecision"].as<string>()
|
|
: "";
|
|
|
|
if (!get_recall_precision(recall_precision_string)) {
|
|
LOG_ERROR("Get recall precision failed, value: %s",
|
|
recall_precision_string.c_str());
|
|
return -1;
|
|
}
|
|
|
|
RetrievalMode retrieval_mode{RM_DENSE};
|
|
if (config_common["RetrievalMode"]) {
|
|
std::string retrieval_mode_str =
|
|
config_common["RetrievalMode"].as<string>();
|
|
if (retrieval_mode_str == "dense") {
|
|
retrieval_mode = RM_DENSE;
|
|
} else if (retrieval_mode_str == "sparse") {
|
|
retrieval_mode = RM_SPARSE;
|
|
}
|
|
}
|
|
|
|
FilterMode filter_mode{FM_NONE};
|
|
if (config_common["FilterMode"]) {
|
|
std::string filter_mode_str = config_common["FilterMode"].as<string>();
|
|
if (filter_mode_str == "tag") {
|
|
filter_mode = FM_TAG;
|
|
}
|
|
}
|
|
|
|
string query_file = config_common["QueryFile"].as<string>();
|
|
|
|
string first_sep = config_common["QueryFirstSep"]
|
|
? config_common["QueryFirstSep"].as<string>()
|
|
: ";";
|
|
string second_sep = config_common["QuerySecondSep"]
|
|
? config_common["QuerySecondSep"].as<string>()
|
|
: " ";
|
|
string query_type = config_common["QueryType"]
|
|
? config_common["QueryType"].as<string>()
|
|
: "float";
|
|
string container_type = config_common["ContainerType"]
|
|
? config_common["ContainerType"].as<string>()
|
|
: "MMapFileStorage";
|
|
|
|
string ground_truth_file = "";
|
|
string ground_truth_first_sep = ";";
|
|
string ground_truth_second_sep = " ";
|
|
|
|
if (config_common["GroundTruthFile"]) {
|
|
ground_truth_file = config_common["GroundTruthFile"].as<string>();
|
|
|
|
if (config_common["GroundTruthFirstSep"]) {
|
|
ground_truth_first_sep =
|
|
config_common["GroundTruthFirstSep"].as<string>();
|
|
}
|
|
|
|
if (config_common["GroundTruthSecondSep"]) {
|
|
ground_truth_second_sep =
|
|
config_common["GroundTruthSecondSep"].as<string>();
|
|
}
|
|
}
|
|
|
|
string index_dir = config_common["IndexPath"].as<string>();
|
|
|
|
core_interface::Index::Pointer index;
|
|
core_interface::BaseIndexQueryParam::Pointer query_param;
|
|
if (parse_and_load_index_param(config_node, index_dir, index, query_param) !=
|
|
0) {
|
|
LOG_ERROR("Failed to parse and load index param");
|
|
return -1;
|
|
}
|
|
|
|
if (retrieval_mode == RM_DENSE) {
|
|
recall_dense(query_type, thread_count, batch_count, top_k, gt_count,
|
|
query_file, first_sep, second_sep, ground_truth_file,
|
|
ground_truth_first_sep, ground_truth_second_sep, index,
|
|
query_param, index_dir, log_dir, filter_mode);
|
|
} else if (retrieval_mode == RM_SPARSE) {
|
|
recall_sparse(query_type, thread_count, batch_count, top_k, gt_count,
|
|
query_file, first_sep, second_sep, ground_truth_file,
|
|
ground_truth_first_sep, ground_truth_second_sep, index,
|
|
query_param, index_dir, log_dir, filter_mode);
|
|
} else {
|
|
std::string mode = retrieval_mode == 1 ? "Dense" : "Sparse";
|
|
LOG_ERROR("unsupported retrieval mode: %s", mode.c_str());
|
|
return -1;
|
|
}
|
|
|
|
// Cleanup
|
|
index->Close();
|
|
|
|
cout << "Recall done." << endl;
|
|
|
|
return 0;
|
|
}
|